defmodule Oban.Pro.Chunk do
  @moduledoc """
  Chunk workers execute jobs together in groups based on a size or a timeout option, e.g. when
  1000 jobs are available or after 10 minutes have elapsed.

  Multiple chunks can run in parallel within a single queue, and each chunk may be composed of
  many thousands of jobs. Aside from improved throughput for a single queue, chunks are ideal as
  the initial stage of data-ingestion and data-processing pipelines.

  ## Usage

  Let's define a worker that sends SMS messages in chunks, rather than individually:

  ```elixir
  defmodule MyApp.ChunkWorker do
    use Oban.Pro.Chunk, queue: :messages, size: 100, timeout: 1000

    @impl true
    def process([_ | _] = jobs) do
      jobs
      |> Enum.map(& &1.args)
      |> MyApp.Messaging.send_batch()

      :ok
    end
  end
  ```

  Notice that we declared a `size` and a `timeout` along with a `queue` for the worker. If `size`
  or `timeout` are omitted they fall back to their defaults: 100 `size` and 1000ms respectively.

  To process _larger_ batches _less_ frequently, we can increase both values:

  ```elixir
  use Oban.Pro.Chunk, size: 500, timeout: :timer.seconds(5)
  ```

  Now chunks will run with up to 500 jobs or every 5 seconds, whichever comes first.

  Unlike other Pro workers, a `Chunk` worker's `process/1` receives a list of jobs rather than a
  single job struct. Fittingly, the [expected return values](#t:result/0) are different as well.

  ## Chunk Partitioning

  By default, chunks are divided into groups based on the `queue` and `worker`. This means that
  each chunk consists of workers belonging to the same queue, regardless of their `args` or
  `meta`. However, this approach may not always be suitable. For instance, you may want to group
  workers based on a specific argument such as `:account_id` instead of just the worker. In such
  cases, you can use the [`:by`](#t:chunk_by/0) option to customize how chunks are partitioned.

  Here are a few examples of the `:by` option that you can use to achieve fine-grained control over
  chunk partitioning:

  ```elixir
  # Explicitly chunk by :worker
  use Oban.Pro.Chunk, by: :worker

  # Chunk by a single args key without considering the worker
  use Oban.Pro.Chunk, by: [args: :account_id]

  # Chunk by multiple args keys without considering the worker
  use Oban.Pro.Chunk, by: [args: [:account_id, :cohort]]

  # Chunk by worker and a single args key
  use Oban.Pro.Chunk, by: [:worker, args: :account_id]

  # Chunk by a single meta key
  use Oban.Pro.Chunk, by: [meta: :source_id]
  ```

  When partitioning chunks of jobs, it's important to note that using only `:args` or `:meta`
  without `:worker` may result in heterogeneous chunks of jobs from different workers.
  Nevertheless, regardless of the partitioning method chunks always consist of jobs from the same
  queue.

  Here's a simple example of partitioning by `worker` and an `author_id` field to batch analysis
  of recent messages per author:

  ```elixir
  defmodule MyApp.LLMAnalysis do
    use Oban.Pro.Chunk, by: [:worker, args: :author_id], size: 50, timeout: 30_000

    @impl true
    def process([%{args: %{"author_id" => author_id}} | _] = jobs) do
      messages =
        jobs
        |> Enum.map(& &1.args["message_id"])
        |> MyApp.Messages.all()

      {:ok, sentiment} = MyApp.GPT.gauge_sentiment(messages)

      MyApp.Messages.record_sentiment(author_id)
    end
  end
  ```

  ## Chunk Results and Error Handling

  `Chunk` worker's result type is tailored to handling multiple jobs at once. The result types
  allow you to succeed an entire chunk, or selectively apply different outcomes to individual
  jobs. Here are the possible results:

  * `:ok` — The chunk succeeded and all jobs can be marked complete

  * `{:ok, result}` — Like `:ok`, but with a result for testing.

  * `{:error, reason, jobs}` — One or more jobs in the chunk failed and may be retried, any
    unlisted jobs are a success.

  * `{:cancel, reason, jobs}` — One or more jobs in the chunk should be cancelled, any unlisted
    jobs are a success.

  * `{:snooze, period, jobs}` — One or more jobs in the chunk should be snoozed and rescheduled
    to run again after the specified period (e.g. `10` for seconds or `{1, :minute}`), any
    unlisted jobs are a success.

  * `[error: {reason, jobs}, cancel: {reason, jobs}, ...]` — Selectively apply different
    outcomes to different jobs. Supports `error`, `cancel`, `discard`, and `snooze` keys. For
    snooze, use `snooze: {period, jobs}`. Any unlisted jobs are a success.

  To illustrate using chunk results let's expand on the message processing example from earlier.
  We'll extend it to complete the whole batch when all messages are delivered or cancel
  undeliverable messages:

  ```elixir
  def process([_ | _] = jobs) do
    notifications =
      jobs
      |> Enum.map(& &1.args)
      |> MyApp.Messaging.send_batch()

    bad_token = fn %{response: response} -> response == :bad_token end

    if Enum.any?(notifications, bad_token) do
      cancels =
        notifications
        |> Enum.zip(jobs)
        |> Enum.filter(fn {notification, _job} -> bad_token.(notification) end)
        |> Enum.map(&elem(&1, 1))

      {:cancel, :bad_token, cancels}
    else
      {:ok, notifications}
    end
  end
  ```

  In the event of an ephemeral crash, like a network issue, the entire chunk may be retried if
  there are any remaining attempts.

  ## Chunk Organization

  Chunks are run by a leader job (which has nothing to do with peer leadership). When the leader
  executes it determines whether a complete chunk is available or if enough time has elapsed to
  run anyhow. If neither case applies, then the leader will sleep until the timeout elapses and
  then execute with as many jobs as it can find.

  Completion queries run every `1000ms` by default. You can use the `:sleep` option to control how
  long the leader delays between queries to check for complete chunks:

  ```elixir
  use Oban.Pro.Chunk, size: 50, sleep: 2_000, timeout: 10_000
  ```

  ## Run Chunks Immediately

  A chunk won't run until either the full number of jobs are available (`size`), or the total time
  has elapsed (`timeout`). In situations where there is a long gap between incoming jobs it may be
  desirable to run immediately when the leader starts. This can be done with the `leading` option.
  For example, to chunk all jobs that arrive within a 5 minute window, but run immediately if more
  than 5 minutes have passed since the last chunk completed:

  ```elixir
  use Oban.Pro.Chunk, size: 100, leading: true, timeout: :timer.minutes(5)
  ```

  ## Optimizing Chunks

  Queues with high concurrency and low throughput may have multiple chunk leaders running
  simultaneously rather than getting bundled into a single chunk. The solution is to reduce the
  queue's global concurrency to a smaller number so that new chunk leader jobs don't start and the
  existing chunk can run a bigger batch.

  ```elixir
  queues: [
    chunked_queue: [global_limit: 1]
  ]
  ```

  A low global limit will reduce overall throughput. Partitioning the queue with the same options
  as the chunk can preserve throughput by allowing one concurrent leader per chunk. For example,
  if you were chunking by `:user_id` in args, you'd partition the queue with the same option:

  ```elixir
  queues: [
    chunked_queue: [local_limit: 20, global_limit: [allowed: 1, partition: [args: :user_id]]]
  ]
  ```
  """

  import Ecto.Query
  import DateTime, only: [utc_now: 0]

  alias Oban.{CrashError, Job, Notifier, PerformError, Period, Repo, Validation}
  alias Oban.Pro.Engines.Smart
  alias Oban.Pro.Utils

  @typedoc """
  A single key or list of keys used to specify fields for partitioning.
  """
  @type key_or_keys :: atom() | [atom()]

  @typedoc """
  Specifies how jobs are grouped into chunks.

  Jobs are always partitioned by queue, but may also be grouped by worker, args, or meta fields.
  """
  @type chunk_by ::
          :worker
          | {:args, key_or_keys()}
          | {:meta, key_or_keys()}
          | [:worker | {:args, key_or_keys()} | {:meta, key_or_keys()}]

  @typedoc """
  A sub-result for error, cancel, or discard outcomes within a keyword list result.
  """
  @type sub_result :: {reason :: term(), [Job.t()]}

  @typedoc """
  A sub-result for snooze outcomes within a keyword list result.
  """
  @type snooze_result :: {period :: Oban.Period.t(), [Job.t()]}

  @typedoc """
  The return value of `process/1`, indicating the outcome for jobs in the chunk.

  Unlisted jobs are marked as completed. See the "Chunk Results and Error Handling" section for
  details on each variant.
  """
  @type result ::
          :ok
          | {:ok, term()}
          | {:cancel | :discard | :error, reason :: term(), [Job.t()]}
          | {:snooze, period :: Oban.Period.t(), [Job.t()]}
          | [
              cancel: sub_result(),
              discard: sub_result(),
              error: sub_result(),
              snooze: snooze_result()
            ]

  @typedoc """
  Options for configuring chunk behavior.
  """
  @type options :: [
          by: chunk_by(),
          leading: boolean(),
          size: pos_integer(),
          sleep: pos_integer(),
          timeout: pos_integer()
        ]

  @doc false
  defmacro __using__(opts) do
    {chunk_opts, other_opts} = Keyword.split(opts, ~w(by leading size sleep timeout)a)

    quote do
      Validation.validate!(unquote(chunk_opts), &Oban.Pro.Chunk.validate/1)

      use Oban.Pro.Worker, unquote(other_opts)

      alias Oban.Pro.Chunk

      @chunk_meta %{
        chunk: true,
        chunk_by: Keyword.get(unquote(chunk_opts), :by, [:worker]),
        chunk_leading: Keyword.get(unquote(chunk_opts), :leading, false),
        chunk_size: Keyword.get(unquote(chunk_opts), :size, 100),
        chunk_sleep: Keyword.get(unquote(chunk_opts), :sleep, 1000),
        chunk_timeout: Keyword.get(unquote(chunk_opts), :timeout, 1000)
      }

      @impl Oban.Worker
      def new(args, opts) when is_map(args) and is_list(opts) do
        meta =
          @chunk_meta
          |> Map.merge(Keyword.get(opts, :meta, %{}))
          |> Map.update!(:chunk_by, &Utils.normalize_by/1)

        chunk_id = Chunk.to_chunk_id(meta.chunk_by, __MODULE__, args, opts)
        meta = Map.put(meta, :chunk_id, chunk_id)

        super(args, Keyword.put(opts, :meta, meta))
      end

      @impl Oban.Worker
      def perform(%Job{} = job) do
        opts = __opts__()

        with {:ok, job} <- Oban.Pro.Worker.before_process(__MODULE__, job, opts) do
          Process.put(:oban_processing, {__MODULE__, job, opts})

          job
          |> Chunk.maybe_process(__MODULE__)
          |> Oban.Pro.Worker.after_process(__MODULE__, job, opts)
        end
      end
    end
  end

  defmacrop in_chunk(meta, chunk_id) do
    quote do
      fragment(
        "? \\? 'chunk_id' AND ?->>'chunk_id' = ?",
        unquote(meta),
        unquote(meta),
        unquote(chunk_id)
      )
    end
  end

  defmacrop finished_at(job) do
    quote do
      type(
        fragment(
          "greatest(?, ?, ?)",
          unquote(job).completed_at,
          unquote(job).discarded_at,
          unquote(job).cancelled_at
        ),
        :utc_datetime_usec
      )
    end
  end

  @doc false
  def validate(opts) do
    Validation.validate_schema(opts,
      by: {:custom, &Oban.Pro.Validation.validate_by/1},
      leading: :boolean,
      size: :pos_integer,
      sleep: :timeout,
      timeout: :timeout
    )
  end

  @doc false
  def to_chunk_id(chunk_by, worker, args, opts) do
    chunk_data =
      Enum.map([:queue | chunk_by], fn
        :queue -> opts |> Keyword.get(:queue, "default") |> to_string()
        :worker -> inspect(worker)
        [:args, keys] -> Utils.take_keys(args, keys)
        [:meta, keys] -> opts |> Keyword.get(:meta, %{}) |> Utils.take_keys(keys)
      end)

    Utils.hash64(chunk_data)
  end

  # Public Interface

  @doc false
  def maybe_process(%Job{conf: conf} = job, worker) do
    cond do
      full_timeout?(conf, job) ->
        fetch_and_process(worker, job)

      full_chunk?(conf, job) ->
        fetch_and_process(worker, job)

      true ->
        job.meta
        |> Map.get("chunk_timeout", 1000)
        |> sleep_and_process(worker, job)
    end
  end

  # Check Helpers

  defp full_chunk?(conf, %{meta: %{"chunk_id" => chunk_id, "chunk_size" => size}}) do
    limit = size - 1

    query =
      Job
      |> select([j], j.id)
      |> where([j], j.state == "available")
      |> where([j], in_chunk(j.meta, ^chunk_id))
      |> limit(^limit)

    Repo.one(conf, from(oj in subquery(query), select: count(oj.id) >= ^limit))
  end

  defp full_timeout?(_, %{meta: %{"chunk_leading" => false}}), do: false

  defp full_timeout?(conf, job) do
    %{id: job_id, meta: %{"chunk_id" => chunk_id, "chunk_timeout" => timeout}} = job

    query =
      Job
      |> where([j], j.id != ^job_id)
      |> where([j], j.state in ~w(completed cancelled discarded))
      |> where([j], in_chunk(j.meta, ^chunk_id))
      |> order_by(desc: :id)
      |> limit(1)
      |> select([j], finished_at(j))

    last_ts = Repo.one(conf, query)

    last_ts &&
      DateTime.utc_now()
      |> DateTime.add(-timeout, :millisecond)
      |> DateTime.after?(last_ts)
  end

  # Processing Helpers

  defp sleep_and_process(total, worker, %Job{conf: conf, meta: meta} = job) do
    sleep = Map.get(meta, "chunk_sleep", 1000)

    Process.sleep(sleep)

    if total - sleep < 0 or full_chunk?(conf, job) do
      fetch_and_process(worker, job)
    else
      sleep_and_process(total - sleep, worker, job)
    end
  end

  defp fetch_and_process(worker, %{conf: conf, meta: %{"chunk_size" => size}} = job) do
    {:ok, chunk} = fetch_chunk(conf, job, size - 1)
    {chunk, _errored} = prepare_chunk(chunk)

    guard_cancel(conf, worker, job, chunk)
    guard_timeout(conf, worker.timeout(job), worker, job, chunk)

    try do
      case worker.process([job | chunk]) do
        :ok ->
          ack_jobs(conf, worker, job, chunk, %{})

        {:ok, result} ->
          with :ok <- ack_jobs(conf, worker, job, chunk, %{}), do: {:ok, result}

        {:cancel, reason, cancelled} ->
          ops = Map.new(cancelled, &{&1.id, {:cancel, reason}})
          ack_jobs(conf, worker, job, chunk, ops)

        {:discard, reason, discarded} ->
          ops = Map.new(discarded, &{&1.id, {:discard, reason}})
          ack_jobs(conf, worker, job, chunk, ops)

        {:error, reason, errored} ->
          ops = Map.new(errored, &{&1.id, {:error, reason}})
          ack_jobs(conf, worker, job, chunk, ops)

        {:snooze, seconds, snoozed} ->
          ops = Map.new(snoozed, &{&1.id, {:snooze, seconds}})
          ack_jobs(conf, worker, job, chunk, ops)

        [_ | _] = multiple ->
          ops =
            for {oper, {reason, jobs}} <- multiple,
                job <- jobs,
                into: %{},
                do: {job.id, {oper, reason}}

          ack_jobs(conf, worker, job, chunk, ops)
      end
    rescue
      error ->
        ops = Map.new(chunk, &{&1.id, {:raised, error}})
        ack_jobs(conf, worker, job, chunk, ops, __STACKTRACE__)

        reraise error, __STACKTRACE__
    catch
      kind, reason ->
        error = CrashError.exception({kind, reason, __STACKTRACE__})
        ops = Map.new(chunk, &{&1.id, {:raised, error}})
        ack_jobs(conf, worker, job, chunk, ops, __STACKTRACE__)

        reraise error, __STACKTRACE__
    end
  end

  # This replicates the query used in Smart.fetch_jobs/3, without the meta tracking or any other
  # complications. Any modifications to the original query must be replicated here. Another
  defp fetch_chunk(conf, %{meta: %{"chunk_id" => chunk_id}} = job, limit) do
    subset_query =
      Job
      |> select([:id])
      |> where([j], j.state == "available")
      |> where([j], in_chunk(j.meta, ^chunk_id))
      |> order_by(asc: :priority, asc: :scheduled_at, asc: :id)
      |> limit(^limit)
      |> lock("FOR UPDATE SKIP LOCKED")

    query =
      Job
      |> with_cte("subset", as: ^subset_query)
      |> join(:inner, [j], x in fragment(~s("subset")), on: true)
      |> where([j, x], j.id == x.id)
      |> select([j, _], j)

    attempted_by = job.attempted_by ++ ["chunk-#{job.id}"]

    updates = [
      set: [state: "executing", attempted_at: DateTime.utc_now(), attempted_by: attempted_by],
      inc: [attempt: 1]
    ]

    Repo.transaction(conf, fn ->
      {_count, chunk} = Repo.update_all(conf, query, updates)

      chunk
    end)
  end

  defp prepare_chunk(chunk) do
    Enum.reduce(chunk, {[], []}, fn job, {acc, err} ->
      with {:ok, worker} <- Oban.Worker.from_string(job.worker),
           {:ok, job} <- Oban.Pro.Worker.before_hook(worker, job, worker.__opts__()) do
        {[job | acc], err}
      else
        {:error, reason} ->
          {acc, [{job, reason} | err]}
      end
    end)
  end

  # Guards

  # Only the leader job is registered as "running" by the queue producer. We listen for pkill
  # messages for _other_ jobs in the chunk and apply that to all jobs. Without this, cancelling
  # doesn't apply to the leader and chunk jobs are orphaned.
  defp guard_cancel(conf, worker, job, chunk) do
    # No supervision tree is running when draining, attempting to notify will raise an exception.
    if Process.get(:oban_draining) do
      :ok
    else
      parent = self()

      Task.start(fn ->
        ref = Process.monitor(parent)
        :ok = Notifier.listen(conf.name, [:signal])

        await_cancel(conf, ref, worker, job, chunk)
      end)
    end
  end

  defp await_cancel(conf, ref, worker, job, chunk) do
    receive do
      {:notification, :signal, %{"action" => "pkill", "job_id" => kill_id}} ->
        if Enum.any?(chunk, &(&1.id == kill_id)) do
          ops = Map.new(chunk, &{&1.id, {:cancel, :shutdown}})

          ack_jobs(conf, worker, job, chunk, ops)

          Notifier.notify(conf.name, :signal, %{action: "pkill", job_id: job.id})
        else
          await_cancel(conf, ref, worker, job, chunk)
        end

      {:DOWN, ^ref, :process, _pid, _reason} ->
        :ok
    end
  end

  # The executor uses :timer.exit_after/2 to kill jobs that exceed the timeout. The queue's
  # producer then catches the DOWN message and uses that to record a job error. The producer
  # isn't aware of the job's chunk, so we monitor the parent process and ack the chunk jobs.
  defp guard_timeout(_conf, :infinity, _worker, _job, _chunk), do: :ok

  defp guard_timeout(conf, timeout, worker, job, chunk) do
    parent = self()

    Task.start(fn ->
      ref = Process.monitor(parent)

      receive do
        {:DOWN, ^ref, :process, _pid, %Oban.TimeoutError{} = error} ->
          ops = Map.new(chunk, &{&1.id, {:raised, error}})

          ack_jobs(conf, worker, job, chunk, ops)

        {:DOWN, ^ref, :process, _pid, _reason} ->
          :ok
      after
        timeout + 100 -> :ok
      end
    end)
  end

  # Acking

  defp ack_jobs(conf, worker, host, chunk, ops, stacktrace \\ []) do
    now = DateTime.utc_now()

    to_ack = fn job ->
      set =
        case Map.get(ops, job.id) do
          nil ->
            %{state: "completed", completed_at: now}

          {:cancel, reason} ->
            error = format_error(job, worker, {:cancel, reason}, [])
            %{state: "cancelled", cancelled_at: now, error: error}

          {:discard, reason} ->
            error = format_error(job, worker, {:discard, reason}, [])
            %{state: "discarded", discarded_at: now, error: error}

          {:error, reason} when job.attempt >= job.max_attempts ->
            error = format_error(job, worker, {:discard, reason}, [])
            %{state: "discarded", discarded_at: now, error: error}

          {:error, reason} ->
            error = format_error(job, worker, {:error, reason}, [])
            %{state: "retryable", scheduled_at: backoff_at(worker, job), error: error}

          {:raised, error} when job.attempt >= job.max_attempts ->
            error = format_error(job, worker, error, stacktrace)
            %{state: "discarded", discarded_at: now, error: error}

          {:raised, error} ->
            error = format_error(job, worker, error, stacktrace)
            %{state: "retryable", scheduled_at: backoff_at(worker, job), error: error}

          {:snooze, period} ->
            scheduled_at = DateTime.add(now, Period.to_seconds(period), :second)
            %{state: "scheduled", scheduled_at: scheduled_at}
        end

      {{:ack, conf.name, job.queue, job.id}, nil, [], set}
    end

    acks =
      chunk
      |> Enum.reject(&(&1.id == host.id))
      |> Enum.map(to_ack)

    with {:ok, _} <- Smart.ack_jobs(acks, conf), do: Map.get(ops, host.id, :ok)
  end

  defp backoff_at(worker, job) do
    DateTime.add(utc_now(), worker.backoff(job), :second)
  end

  defp format_error(job, worker, reason, stacktrace) when is_tuple(reason) do
    format_error(job, worker, PerformError.exception({worker, reason}), stacktrace)
  end

  defp format_error(job, _worker, error, stacktrace) do
    {blamed, stacktrace} = Exception.blame(:error, error, stacktrace)

    formatted = Exception.format(:error, blamed, stacktrace)

    %{attempt: job.attempt, at: utc_now(), error: formatted}
  end
end