defmodule Oban.Pro.Plugins.DynamicPrioritizer do
  @moduledoc """
  The DynamicPrioritizer plugin automatically adjusts job's 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.

  ## 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. To reprioritize after less time waiting you can
  configure the `:after` time:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicPrioritizer, after: :timer.minutes(2)}]
  ```

  Now lower job priorities are bumped after 2 minutes of waiting, and every minute after that. To
  lower the reprioritization frequency you can change the `:interval` along with the `:after`
  time:

  ```elixir
  plugins: [{
    Oban.Pro.Plugins.DynamicPrioritizer,
    after: :timer.minutes(10),
    interval: :timer.minutes(5)
  }]
  ```

  Here we've specified that job priority will increase every 5 minutes after the first 10 minutes
  of waiting.

  It's also possible to limit the maximum priority jobs may transition to. Typically, jobs will
  converge on `0` priority, as that's the highest priority. Set the `max_priority` option to
  change to a lower priority:

  ```elixir
  plugins: [{Oban.Pro.Plugins.DynamicPrioritizer, max_priority: 2}]
  ```

  ## Providing Overrides

  The `after` option applies to jobs for any workers across all queues. The `DynamicPrioritizer`
  plugin allows you to specify per-queue and per-worker overrides that fine tune reprioritization.

  Let's configure reprioritization for the `analysis` queue so that it nudges jobs after only 1
  minute:

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

  We can also effectively disable reprioritization for all other queues by setting the period to
  `:infinity`:

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

  If per-queue overrides aren't enough we can override on a per-worker basis instead:

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

  Naturally you can mix and match overrides to finely control reprioritization:

  ```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
