defmodule Oban.Pro.Plugins.DynamicLifeline do
  @moduledoc """
  The `DynamicLifeline` plugin uses producer records to periodically rescue orphaned jobs, i.e.
  jobs that are stuck in the `executing` state because the node was shut down before the job could
  finish. In addition, it performs the following maintenance tasks:

  * Discard jobs left `available` with an attempt equal to max attempts
  * Repair stuck workflows with deleted dependencies or missed scheduling events
  * Repair stuck chains with deleted dependencies or missed scheduling events
  * Repair jobs in partitioned queues that are missing a partition key

  Without `DynamicLifeline` you'll need to manually rescue stuck jobs or perform maintenance.

  ## Using the Plugin

  To use the `DynamicLifeline` plugin, add the module to your list of Oban plugins in
  `config.exs`:

  ```elixir
  config :my_app, Oban,
    plugins: [Oban.Pro.Plugins.DynamicLifeline]
    ...
  ```

  There isn't any configuration necessary. By default, the plugin rescues orphaned jobs every 1
  minute. If necessary, you can override the rescue interval:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicLifeline, rescue_interval: :timer.minutes(5)}]
  ```

  If your system is under high load or produces a multitude of orphans you may wish to increase
  the query timeout beyond the `45s` default:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicLifeline, timeout: :timer.minutes(1)}]
  ```

  Note that rescuing orphans relies on producer records as used by the `Smart` engine.

  ## Repairing Workflows, Chains, and Partitions

  The plugin automatically repairs stuck workflows, chains, and partitioned jobs during each
  rescue cycle. These repair operations ensure:

  * **Workflow repair** — Jobs held waiting for dependencies that were deleted or missed a
    scheduling event are released. This handles edge cases where workflow jobs get stuck due to
    incomplete dependency resolution.

  * **Chain repair** — Similar to workflows, chain jobs waiting on deleted or stuck predecessors
    are released to continue processing.

  * **Partition repair** — Jobs in partitioned queues that are missing a `partition_key` in their
    metadata (e.g., jobs scheduled before a queue was partitioned) have their partition key
    computed and set automatically.

  By default, each repair operation processes up to 1000 jobs per cycle. You can adjust this
  limit with the `repair_limit` option:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicLifeline, repair_limit: 500}]
  ```

  ## Identifying Rescued Jobs

  Rescued jobs can be identified by a `rescued` value in `meta`. Each rescue increments the
  `rescued` count by one.

  ## Rescuing Exhausted Jobs

  When a job's `attempt` matches its `max_attempts` its retries are considered "exhausted".
  Normally, the `DynamicLifeline` plugin transitions exhausted jobs to the `discarded` state and
  they won't be retried again. It does this for a couple of reasons:

  1. To ensure at-most-once semantics. Suppose a long-running job interacted with a non-idempotent
     service and was shut down while waiting for a reply; you may not want that job to retry.

  2. To prevent infinitely crashing BEAM nodes. Poorly behaving jobs may crash the node (through
     NIFs, memory exhaustion, etc.) We don't want to repeatedly rescue and rerun a job that
     repeatedly crashes the entire node.

  Discarding exhausted jobs may not always be desired. Use the `retry_exhausted` option if you'd
  prefer to retry exhausted jobs when they are rescued, rather than discarding them:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicLifeline, retry_exhausted: true}]
  ```

  During rescues, with `retry_exhausted: true`, a job's `max_attempts` is incremented and it is
  moved back to the `available` state.

  ## Instrumenting with Telemetry

  The `DynamicLifeline` plugin adds the following metadata to the `[:oban, :plugin, :stop]` event:

  * `:rescued_jobs` — a list of jobs transitioned back to `available`

  * `:discarded_jobs` — a list of jobs transitioned to `discarded`

  _Note: jobs only include `id`, `queue`, and `state` fields._
  """

  @behaviour Oban.Plugin

  use GenServer

  import Ecto.Query

  alias Oban.Pro.Engines.Smart
  alias Oban.Pro.Stages.Chain
  alias Oban.Pro.{Partition, Producer, Workflow}
  alias Oban.{Job, Peer, Repo, Validation}

  require Logger

  @type option ::
          {:conf, Oban.Config.t()}
          | {:name, Oban.name()}
          | {:repair_limit, pos_integer()}
          | {:retry_exhausted, boolean()}
          | {:rescue_interval, timeout()}
          | {:timeout, timeout()}

  defstruct [
    :conf,
    :rescue_timer,
    repair_limit: 1000,
    retry_exhausted: false,
    rescue_interval: :timer.minutes(1),
    timeout: :timer.seconds(45)
  ]

  defmacrop attempted_join(attempted_by, uuid) do
    quote do
      fragment(
        "array_length(?, 1) <> 1 and ? = uuid (?[2])",
        unquote(attempted_by),
        unquote(uuid),
        unquote(attempted_by)
      )
    end
  end

  defmacrop coalesce_partition(meta) do
    quote do
      fragment(
        "COALESCE(?->'global_limit'->'partition', ?->'rate_limit'->'partition')",
        unquote(meta),
        unquote(meta)
      )
    end
  end

  defmacrop track_rescued(meta) do
    quote do
      fragment(
        """
        jsonb_set(?, '{rescued}', (COALESCE(?->>'rescued', '0')::int + 1)::text::jsonb)
        """,
        unquote(meta),
        unquote(meta)
      )
    end
  end

  @doc false
  def child_spec(args), do: super(args)

  @impl Oban.Plugin
  @spec start_link([option()]) :: GenServer.on_start()
  def start_link(opts) do
    {name, opts} = Keyword.pop(opts, :name)

    case opts[:conf] do
      %{engine: Smart} ->
        GenServer.start_link(__MODULE__, struct!(__MODULE__, opts), name: name)

      %{engine: engine} ->
        raise RuntimeError, """
        DynamicLifeline requires the Smart engine to run correctly, but you're using:

        engine: #{inspect(engine)}

        You can either switch to use the Smart or remove DynamicLifeline from your plugins.
        """
    end
  end

  @impl Oban.Plugin
  def validate(opts) do
    Validation.validate_schema(opts,
      conf: :any,
      name: :any,
      repair_limit: :pos_integer,
      rescue_interval: :pos_integer,
      rescue_after: :pos_integer,
      retry_exhausted: :boolean,
      timeout: :timeout
    )
  end

  @impl Oban.Plugin
  def format_logger_output(_conf, meta) do
    for {key, val} <- meta,
        key in ~w(discarded_jobs rescued_jobs)a,
        into: %{},
        do: {key, Enum.map(val, & &1.id)}
  end

  @impl GenServer
  def init(state) do
    Process.flag(:trap_exit, true)

    :telemetry.execute([:oban, :plugin, :init], %{}, %{conf: state.conf, plugin: __MODULE__})

    {:ok, schedule_rescue(state)}
  end

  @impl GenServer
  def terminate(_reason, state) do
    if is_reference(state.rescue_timer), do: Process.cancel_timer(state.rescue_timer)

    :ok
  end

  @impl GenServer
  def handle_info(:rescue, state) do
    if Peer.leader?(state.conf) do
      meta = %{conf: state.conf, plugin: __MODULE__}

      :telemetry.span([:oban, :plugin], meta, fn ->
        orphan_meta = rescue_orphaned(state)

        :ok = repair_workflows(state)
        :ok = repair_chains(state)
        :ok = repair_partitions(state)

        {:ok, Map.merge(meta, orphan_meta)}
      end)
    end

    {:noreply, schedule_rescue(state)}
  end

  def handle_info(message, state) do
    Logger.warning(
      message: "Received unexpected message: #{inspect(message)}",
      source: :oban_pro,
      module: __MODULE__
    )

    {:noreply, state}
  end

  # Scheduling

  defp schedule_rescue(state) do
    timer = Process.send_after(self(), :rescue, state.rescue_interval)

    %{state | rescue_timer: timer}
  end

  # Orphan Rescuing

  defp rescue_orphaned(state) do
    {res_count, res_jobs} = transition_available(orphan_query(), state)
    {dis_count, dis_jobs} = transition_discarded(orphan_query(), state)
    {exh_count, exh_jobs} = transition_discarded(stuck_query(), state)

    if state.retry_exhausted do
      %{
        discarded_count: 0,
        discarded_jobs: [],
        rescued_count: res_count + dis_count + exh_count,
        rescued_jobs: res_jobs ++ dis_jobs ++ exh_jobs
      }
    else
      %{
        discarded_count: dis_count + exh_count,
        discarded_jobs: dis_jobs ++ exh_jobs,
        rescued_count: res_count,
        rescued_jobs: res_jobs
      }
    end
  end

  defp orphan_query do
    subquery =
      Job
      |> where([j], not is_nil(j.queue) and j.state == "executing")
      |> join(:left, [j], p in Producer, on: attempted_join(j.attempted_by, p.uuid))
      |> where([_, p], is_nil(p.uuid))

    Job
    |> join(:inner, [j], x in subquery(subquery), on: j.id == x.id)
    |> select([_, x], map(x, [:id, :meta, :queue, :state]))
  end

  # We don't know how it's possible for a job to be available with attempts >= max_attempts, but
  # we can clean up to prevent clogging a queue at least.
  defp stuck_query do
    Job
    |> where([j], j.state == "available")
    |> select([j], map(j, [:id, :meta, :queue, :state]))
  end

  defp transition_available(query, state) do
    query =
      query
      |> where([j], j.attempt < j.max_attempts)
      |> update([j], set: [meta: track_rescued(j.meta), state: "available"])

    Repo.update_all(state.conf, query, [], timeout: state.timeout)
  catch
    :error, %Postgrex.Error{postgres: %{code: :unique_violation, detail: detail}} ->
      Smart.clear_uniq_violation(state.conf, detail)

      {0, []}
  end

  defp transition_discarded(query, %{retry_exhausted: true} = state) do
    query =
      query
      |> where([j], j.attempt >= j.max_attempts)
      |> update([j],
        inc: [max_attempts: 1],
        set: [meta: track_rescued(j.meta), state: "available"]
      )

    Repo.update_all(state.conf, query, [], timeout: state.timeout)
  catch
    :error, %Postgrex.Error{postgres: %{code: :unique_violation, detail: detail}} ->
      Smart.clear_uniq_violation(state.conf, detail)

      {0, []}
  end

  defp transition_discarded(query, state) do
    Repo.update_all(
      state.conf,
      where(query, [j], j.attempt >= j.max_attempts),
      [set: [state: "discarded", discarded_at: DateTime.utc_now()]],
      timeout: state.timeout
    )
  end

  # Composition Repair

  defp repair_workflows(state) do
    Workflow.rescue_workflows(state.conf, timeout: state.timeout, limit: state.repair_limit)
  end

  defp repair_chains(state) do
    Chain.rescue_chains(state.conf, timeout: state.timeout, limit: state.repair_limit)
  end

  # Partition Repair

  defp repair_partitions(state) do
    for {queue, partition} <- partitioned_queues(state) do
      query =
        Job
        |> where([j], j.queue == ^queue)
        |> where([j], j.state in ~w(available scheduled retryable))
        |> where([j], fragment("NOT ? \\? 'partition_key'", j.meta))
        |> limit(^state.repair_limit)

      state.conf
      |> Repo.all(query, timeout: state.timeout)
      |> Enum.group_by(fn job -> gen_partition_key(job, partition) end)
      |> Enum.each(fn {key, grouped_jobs} ->
        ids = Enum.map(grouped_jobs, & &1.id)

        Job
        |> where([j], j.id in ^ids)
        |> update([j], set: [meta: fragment("jsonb_set(?, '{partition_key}', ?)", j.meta, ^key)])
        |> then(&Repo.update_all(state.conf, &1, [], timeout: state.timeout))
      end)
    end

    :ok
  end

  defp gen_partition_key(job, partition) do
    job
    |> Ecto.Changeset.change()
    |> Partition.gen_key(partition)
  end

  defp partitioned_queues(state) do
    normalize_partition = fn partition ->
      fields = partition |> Map.get("fields", []) |> Enum.map(&String.to_existing_atom/1)

      %{fields: fields, keys: Map.get(partition, "keys", [])}
    end

    Producer
    |> where([p], fragment("jsonb_typeof(?) = 'object'", coalesce_partition(p.meta)))
    |> select([p], {p.queue, coalesce_partition(p.meta)})
    |> distinct(true)
    |> then(&Repo.all(state.conf, &1, timeout: state.timeout))
    |> Enum.map(fn {queue, partition} -> {queue, normalize_partition.(partition)} end)
  end
end
