defmodule Oban.Pro.Plugins.DynamicPrioritizer do
  @moduledoc """
  The `DynamicPrioritizer` plugin automatically adjusts job priorities to ensure all jobs are
  eventually processed.

  Using mixed priorities in a queue causes certain jobs to execute before others. For example, a
  queue that processes jobs from various customers may prioritize customers that are in a higher
  tier or plan. All high priority (`0`) jobs are guaranteed to run before any with lower priority
  (`1..9`), which is wonderful for the higher tier customers but can lead to resource starvation.
  When there is a constant flow of high priority jobs the lower priority jobs will never get the
  chance to run.

  The `DynamicPrioritizer` solves this by periodically boosting the priority of jobs that have
  been waiting longer than a configured threshold. Each cycle, qualifying jobs have their priority
  decremented by one (e.g. from `3` to `2`), gradually converging toward the highest priority.

  ## Using the Plugin

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

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

  Without any additional options the plugin will automatically increase the priority of any jobs
  that are `available` for 5 minutes or more, checking once per minute.

  ## Options

  * `:after` — how long a job must be `available` before it becomes eligible for reprioritization,
    in milliseconds. Set to `:infinity` to disable reprioritization by default (useful when
    relying entirely on overrides). Defaults to `300_000` (5 minutes).

  * `:interval` — the number of milliseconds between reprioritization cycles. Defaults to
    `60_000` (1 minute).

  * `:limit` — the maximum number of jobs to reprioritize per override or default group in each
    cycle. Defaults to `10_000`.

  * `:max_priority` — the highest priority (lowest number) that jobs can be boosted to. For
    example, `max_priority: 2` prevents jobs from being boosted above priority `2`, reserving
    `0` and `1` for genuinely high priority work. Defaults to `0`.

  * `:queue_overrides` — a keyword list of per-queue wait thresholds that override `:after` for
    specific queues. See "Queue and Worker Overrides" for details.

  * `:worker_overrides` — a keyword list of per-worker wait thresholds that override `:after` for
    specific workers. See "Queue and Worker Overrides" for details.

  ## Queue and Worker Overrides

  The `:after` threshold applies globally to all queues and workers. Overrides let you fine tune
  reprioritization for specific queues or workers.

  Configure the `analysis` queue to nudge jobs after only 1 minute:

  ```elixir
  plugins: [{
    Oban.Pro.Plugins.DynamicPrioritizer,
    queue_overrides: [analysis: :timer.minutes(1)]
  }]
  ```

  Disable reprioritization globally while enabling it for a single queue:

  ```elixir
  plugins: [{
    Oban.Pro.Plugins.DynamicPrioritizer,
    after: :infinity,
    queue_overrides: [analysis: :timer.minutes(1)]
  }]
  ```

  Override on a per-worker basis:

  ```elixir
  plugins: [{
    Oban.Pro.Plugins.DynamicPrioritizer,
    worker_overrides: [
      "MyApp.HighSLAWorker": :timer.seconds(30),
      "MyApp.LowSLAWorker": :timer.minutes(10)
    ]
  }]
  ```

  Queue and worker overrides can be combined. Note that overrides are applied independently—a
  job matching both a queue override and a worker override will be boosted by each:

  ```elixir
  plugins: [{
    Oban.Pro.Plugins.DynamicPrioritizer,
    interval: :timer.minutes(2),
    after: :timer.minutes(5),
    queue_overrides: [media: :timer.minutes(10)],
    worker_overrides: ["MyApp.HighSLAWorker": :timer.seconds(30)]
  }]
  ```

  ## Instrumenting with Telemetry

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

  * `:reprioritized_count` — the number of jobs reprioritized
  """

  @behaviour Oban.Plugin

  use GenServer

  import Ecto.Query, only: [join: 5, limit: 2, order_by: 2, where: 3]

  alias Oban.{Job, Peer, Repo, Validation}

  require Logger

  @type option :: [
          after: timeout(),
          interval: timeout(),
          limit: pos_integer(),
          max_priority: 0..9,
          queue_overrides: [{atom() | String.t(), timeout()}],
          worker_overrides: [{atom() | String.t(), timeout()}]
        ]

  defstruct [
    :conf,
    :timer,
    after: :timer.minutes(5),
    interval: :timer.seconds(60),
    limit: 10_000,
    max_priority: 0,
    queue_overrides: [],
    worker_overrides: []
  ]

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

  # Callbacks

  @impl Oban.Plugin
  def start_link(opts) do
    {name, opts} = Keyword.pop(opts, :name)

    GenServer.start_link(__MODULE__, struct!(__MODULE__, opts), name: name)
  end

  @impl Oban.Plugin
  def validate(opts) do
    Validation.validate_schema(opts,
      conf: :any,
      name: :any,
      after: :timeout,
      interval: :pos_integer,
      limit: :pos_integer,
      max_priority: {:range, 0..9},
      queue_overrides: {:custom, &validate_overrides(:queues, &1)},
      worker_overrides: {:custom, &validate_overrides(:workers, &1)}
    )
  end

  @impl Oban.Plugin
  def format_logger_output(_conf, meta) do
    Map.take(meta, ~w(reprioritized_count)a)
  end

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

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

    {:ok, schedule_reprioritization(state)}
  end

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

    :ok
  end

  @impl GenServer
  def handle_info(:reprioritize, state) do
    meta = %{conf: state.conf, plugin: __MODULE__}

    :telemetry.span([:oban, :plugin], meta, fn ->
      case reprioritize_starved_jobs(state) do
        {:ok, count} when is_integer(count) ->
          {:ok, Map.put(meta, :reprioritized_count, count)}

        error ->
          {:error, Map.put(meta, :error, error)}
      end
    end)

    {:noreply, schedule_reprioritization(state)}
  end

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

    {:noreply, state}
  end

  # Validation

  defp validate_overrides(parent_key, overrides) do
    Validation.validate(parent_key, overrides, &validate_override/1)
  end

  defp validate_override({key, value}) when is_atom(key) do
    Validation.validate_timeout(key, value)
  end

  defp validate_override(option) do
    {:error, "expected override option to be a tuple, got: #{inspect(option)}"}
  end

  # Scheduling

  defp schedule_reprioritization(state) do
    %{state | timer: Process.send_after(self(), :reprioritize, state.interval)}
  end

  # Queries

  defp reprioritize_starved_jobs(state) do
    if Peer.leader?(state.conf) do
      Repo.transaction(state.conf, fn ->
        queue_counts =
          for {queue, period} <- state.queue_overrides do
            state
            |> base_query()
            |> query_for_queues(:any, [to_string(queue)])
            |> query_for_period(period)
            |> update_all(state)
          end

        worker_counts =
          for {worker, period} <- state.worker_overrides do
            state
            |> base_query()
            |> query_for_workers(:any, [to_string(worker)])
            |> query_for_period(period)
            |> update_all(state)
          end

        default_count =
          state
          |> base_query()
          |> query_for_queues(:not, string_keys(state.queue_overrides))
          |> query_for_workers(:not, string_keys(state.worker_overrides))
          |> query_for_period(state.after)
          |> update_all(state)

        [queue_counts, worker_counts]
        |> List.flatten()
        |> Enum.reduce(default_count, &(&1 + &2))
      end)
    else
      {:ok, 0}
    end
  end

  defp string_keys(keyword) do
    for {key, _} <- keyword, do: to_string(key)
  end

  defp base_query(%{limit: limit, max_priority: max_priority}) do
    Job
    |> where([j], j.state == "available")
    |> where([j], j.priority > ^max_priority)
    |> order_by(asc: :id)
    |> limit(^limit)
  end

  defp query_for_period(query, :infinity) do
    where(query, [j], true == false)
  end

  defp query_for_period(query, period) do
    where(query, [j], j.scheduled_at <= ^to_timestamp(period))
  end

  defp query_for_queues(query, :not, []), do: query
  defp query_for_queues(query, :not, queues), do: where(query, [j], j.queue not in ^queues)
  defp query_for_queues(query, :any, queues), do: where(query, [j], j.queue in ^queues)

  defp query_for_workers(query, :not, []), do: query
  defp query_for_workers(query, :not, workers), do: where(query, [j], j.worker not in ^workers)
  defp query_for_workers(query, :any, workers), do: where(query, [j], j.worker in ^workers)

  defp update_all(subquery, state) do
    query = join(Job, :inner, [j], x in subquery(subquery), on: j.id == x.id)

    state.conf
    |> Repo.update_all(query, inc: [priority: -1])
    |> elem(0)
  end

  defp to_timestamp(milliseconds) do
    DateTime.add(DateTime.utc_now(), -milliseconds, :millisecond)
  end
end
