defmodule Oban.Pro.Batch do
  @moduledoc """
  Batches link the execution of many jobs as a group and runs optional callbacks after jobs are
  processed. This allows your application to coordinate the execution of any number of jobs in
  parallel.

  ## Usage

  Batches are composed of one or more Pro workers that are linked with a shared `batch_id` and
  optional callbacks. As a simple example, let's define a worker that delivers daily emails:

  ```elixir
  defmodule MyApp.EmailBatch do
    use Oban.Pro.Worker, queue: :mailers

    @behaviour Oban.Pro.Batch

    @impl Oban.Pro.Worker
    def process(%Job{args: %{"email" => email}}) do
      MyApp.Mailer.daily_update(email)
    end

    @impl Oban.Pro.Batch
    def batch_completed(_job) do
      Logger.info("BATCH COMPLETE")

      :ok
    end
  end
  ```

  Now, create a batch with `new/1` by passing a list of job changesets:

  ```elixir
  emails
  |> Enum.map(&MyApp.EmailBatch.new(%{email: &1}))
  |> Batch.new()
  |> Oban.insert_all()
  ```

  After all jobs in the batch are `completed`, the `batch_completed/1` callback will be
  triggered.

  ## Handler Callbacks

  After jobs in the batch are processed, a callback job may be inserted. There are four optional
  batch handler callbacks that a worker may define:

  | callback              | enqueued after                                           |
  | --------------------- | -------------------------------------------------------- |
  | `c:batch_attempted/1` | _all_ jobs `executed` at least once                      |
  | `c:batch_cancelled/1` | the _first_ job is `cancelled`                           |
  | `c:batch_completed/1` | _all_ jobs in the batch are `completed`                  |
  | `c:batch_discarded/1` | the _first_ job is `discarded`                           |
  | `c:batch_exhausted/1` | _all_ jobs are `completed`, `cancelled`, or `discarded`  |
  | `c:batch_retryable/1` | the _first_ job is `retryable`                           |

  Each callback runs in a separate, isolated job, so it may be retried or discarded like any other
  job. The callback function receives an `Oban.Job` struct with the `batch_id` in `meta` and
  should return `:ok` (or another valid `Worker.result()`).

  Here we'll implement each of the optional handler callbacks and have them print out the batch
  status along with the `batch_id`:

  ```elixir
  defmodule MyApp.BatchWorker do
    use Oban.Pro.Worker

    @behaviour Oban.Pro.Batch

    @impl Oban.Pro.Worker
    def process(_job), do: :ok

    @impl Oban.Pro.Batch
    def batch_attempted(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:attempted,  batch_id})
      :ok
    end

    @impl Oban.Pro.Batch
    def batch_cancelled(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:cancelled,  batch_id})
      :ok
    end

    @impl Oban.Pro.Batch
    def batch_completed(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:completed,  batch_id})
      :ok
    end

    @impl Oban.Pro.Batch
    def batch_discarded(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:discarded,  batch_id})
      :ok
    end

    @impl Oban.Pro.Batch
    def batch_exhausted(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:exhausted,  batch_id})
      :ok
    end

    @impl Oban.Pro.Batch
    def batch_retryable(%Job{meta: %{"batch_id" => batch_id}}) do
      IO.inspect({:retryable,  batch_id})
      :ok
    end
  end
  ```

  ## Hybrid Batches

  Batches may also be built from a variety of workers, though you must provide an explicit worker
  for callbacks:

  ```elixir
  mail_jobs = Enum.map(mail_args, &MyApp.MailWorker.new/1)
  push_jobs = Enum.map(push_args, &MyApp.PushWorker.new/1)

  [callback_worker: MyApp.CallbackWorker]
  |> Batch.new()
  |> Batch.add(mail_jobs)
  |> Batch.add(push_jobs)
  ```

  Without an explicit `callback_worker`, any worker in the batch may be used for callbacks. That
  makes callback handling unpredictable and unexpected.

  ## Customizing Batches

  Batches accept a variety of options for grouping and callback customization.

  ### Batch ID

  By default a `batch_id` is generated as a time-ordered random [UUIDv7][uuid]. UUIDs are
  sufficient to ensure uniqueness between workers and nodes for any time period. However, if you
  require control, you can override `batch_id` generation at the worker level or pass a value
  directly to the `new/2` function.

  ```elixir
  Batch.new(batch_id: "custom-batch-id")
  ```

  ### Batch Name

  Batches accept an optional name to describe the purpose of the batch, beyond the generated id or
  individual jobs in it. While the `batch_id` must be unique, the `batch_name` doesn't have to be,
  so it can be used as a general purpose label.

  ```elixir
  Batch.new(batch_name: "nightly-etl")
  ```

  ### Callback Workers

  For some batches, notably those with heterogeneous jobs, it's necessary to specify a different
  worker for callbacks. That is easily accomplished by passing the `:callback_worker` option to
  `new/2`:

  ```elixir
  Batch.new(callback_worker: MyCallbackWorker)
  ```

  The callback worker **must** be an `Oban.Pro.Worker` that defines one or more of the batch callback
  handlers.

  ### Callback Options

  By default, callback jobs have an empty `args` map and inherit other options from batch jobs.
  With `callback_opts`, you can set standard job options for batch callback jobs (only `args`,
  `max_attempts`, `meta`, `priority`, `queue`, and `tags` are allowed).

  For example, here we're passing a webhook URLs as `args`:

  ```elixir
  Batch.new(callback_opts: [args: %{webhook: "https://web.hook"}])
  ```

  Here, we change the callback queue and knock the priority down:

  ```elixir
  Batch.new(callback_opts: [queue: :callbacks, priority: 9])
  ```

  Be careful to minimize `callback_opts` as they are stored in each batch job's meta.

  ## Fetching Batch Jobs

  To pull more context from batch jobs, it's possible to load all jobs from the batch with
  `all_jobs/2` and `stream_jobs/2`. The functions take a single batch job and returns a list or
  stream of all non-callback jobs in the batch, which you can then operate on with `Enum` or
  `Stream` functions.

  As an example, imagine you have a batch that ran for a few thousand accounts and you'd like to
  notify admins that the batch is complete.

  ```elixir
  defmodule MyApp.BatchWorker do
    use Oban.Pro.Worker

    @behaviour Oban.Pro.Batch

    @impl Oban.Pro.Batch
    def batch_completed(%Job{} = job) do
      {:ok, account_ids} =
          job
          |> Oban.Pro.Batch.all_jobs()
          |> Enum.map(& &1.args["account_id"])

      account_ids
      |> MyApp.Accounts.all()
      |> MyApp.Mailer.notify_admins_about_batch()
    end
  ```

  Fetching a thousand jobs may be alright, but for larger batches you don't want to load that much
  data into memory. Instead, you can use `stream_jobs` to iterate through them lazily:

  ```elixir
  {:ok, account_ids} =
    MyApp.Repo.transaction(fn ->
      job
      |> Batch.stream_jobs()
      |> Stream.map(& &1.args["account_id"])
      |> Enum.to_list()
    end)
  ```

  Streaming is provided by Ecto's `Repo.stream`, and it must take place within a transaction.
  While it may be overkill for small batches, for batches with tens or hundreds of thousands of
  jobs, it will prevent massive memory spikes or the database grinding to a halt.

  [uuid]: https://datatracker.ietf.org/doc/html/draft-peabody-dispatch-new-uuid-format-04#section-5.2
  """

  @behaviour Oban.Pro.Flusher
  @behaviour Oban.Pro.Handler

  import Ecto.Query, only: [limit: 2, order_by: 2, select: 3, union_all: 2, where: 3]

  alias Ecto.Changeset
  alias Oban.{Job, Repo, Validation, Worker}
  alias Oban.Pro.Engines.Smart
  alias Oban.Pro.Workflow

  require Logger

  @type changeset :: Job.changeset()

  @type callback_opts :: [
          args: Job.args(),
          max_attempts: pos_integer(),
          meta: map(),
          priority: 0..9,
          queue: atom() | String.t(),
          tags: Job.tags()
        ]

  @type batch_opts :: [
          batch_id: String.t(),
          batch_name: String.t(),
          callback_opts: callback_opts(),
          callback_worker: module()
        ]

  @type repo_opts :: [timeout: timeout()]

  @type t :: %__MODULE__{changesets: Enumerable.t(changeset), opts: map()}

  @callbacks_to_functions %{
    "attempted" => :batch_attempted,
    "cancelled" => :batch_cancelled,
    "completed" => :batch_completed,
    "discarded" => :batch_discarded,
    "exhausted" => :batch_exhausted,
    "retryable" => :batch_retryable
  }

  @callbacks_to_deprecated %{
    "attempted" => :handle_attempted,
    "cancelled" => :handle_cancelled,
    "completed" => :handle_completed,
    "discarded" => :handle_discarded,
    "exhausted" => :handle_exhausted,
    "retryable" => :handle_retryable
  }

  @callbacks_to_combined Map.merge(
                           @callbacks_to_functions,
                           @callbacks_to_deprecated,
                           fn _k, new, old -> [new, old] end
                         )

  @callbacks_to_states %{
    "attempted" => ~w(scheduled available executing),
    "completed" => ~w(scheduled available executing retryable cancelled discarded),
    "cancelled" => ~w(cancelled),
    "discarded" => ~w(discarded),
    "exhausted" => ~w(scheduled retryable available executing),
    "retryable" => ~w(retryable)
  }

  @all_states Enum.map(Job.states(), &to_string/1)

  @doc """
  Called after all jobs in the batch were attempted at least once.

  If a `batch_attempted/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when all jobs were attempted:

      def batch_attempted(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Attempted")
      end
  """
  @callback batch_attempted(job :: Job.t()) :: Worker.result()

  @doc """
  Called after any jobs in the batch have a `cancelled` state.

  If a `batch_cancelled/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when any jobs are discarded:

      def batch_cancelled(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Cancelled")
      end
  """
  @callback batch_cancelled(job :: Job.t()) :: Worker.result()

  @doc """
  Called after all jobs in the batch have a `completed` state.

  If a `batch_completed/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when all jobs are completed:

      def batch_completed(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Completed")
      end
  """
  @callback batch_completed(job :: Job.t()) :: Worker.result()

  @doc """
  Called after any jobs in the batch have a `discarded` state.

  If a `batch_discarded/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when any jobs are discarded:

      def batch_discarded(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Discarded")
      end
  """
  @callback batch_discarded(job :: Job.t()) :: Worker.result()

  @doc """
  Called after all jobs in the batch have either a `cancelled`, `completed` or `discarded` state.

  If a `batch_exhausted/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when all jobs are completed or discarded:

      def batch_exhausted(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Exhausted")
      end
  """
  @callback batch_exhausted(job :: Job.t()) :: Worker.result()

  @doc """
  Called after any jobs in the batch have a `retryable` state.

  If a `batch_retryable/1` function is defined it is executed by an isolated callback job.

  ## Example

  Print when any jobs are retryable:

      def batch_retryable(%Job{meta: %{"batch_id" => batch_id}}) do
        IO.inspect(batch_id, label: "Retryable")
      end
  """
  @callback batch_retryable(job :: Job.t()) :: Worker.result()

  @optional_callbacks batch_attempted: 1,
                      batch_cancelled: 1,
                      batch_completed: 1,
                      batch_discarded: 1,
                      batch_exhausted: 1,
                      batch_retryable: 1

  defstruct changesets: [], opts: []

  defguardp is_list_or_stream(enum) when is_list(enum) or is_struct(enum, Stream)

  # Handler Callbacks

  @impl Oban.Pro.Handler
  def on_start do
    events = [
      [:oban, :engine, :cancel_all_jobs, :stop]
    ]

    :telemetry.attach_many("oban.batch", events, &__MODULE__.handle_event/4, nil)
  end

  @impl Oban.Pro.Handler
  def on_stop do
    :telemetry.detach("oban.batch")
  end

  @doc false
  def handle_event(_event, _timing, %{conf: conf, jobs: jobs}, _) do
    for %{meta: %{"batch_id" => batch_id}} = job <- jobs do
      on_flush(job, batch_id, conf)
    end
  end

  # Constants

  @doc false
  def callbacks_to_functions, do: @callbacks_to_functions

  @doc false
  def callbacks_to_deprecated, do: @callbacks_to_deprecated

  # Public Functions

  @doc false
  def new, do: new([], [])

  @doc false
  def new([%Changeset{} | _] = changesets), do: new(changesets, [])
  def new(%Stream{} = changesets), do: new(changesets, [])
  def new(opts) when is_list(opts), do: new([], opts)

  @doc """
  Build a new batch from a list or stream of job changesets and some options.

  ## Examples

  Build an empty batch without any jobs or options:

      Batch.new()

  Build a batch from a list of changesets:

      1..3
      |> Enum.map(fn id -> MyWorker.new(%{id: id}) end)
      |> Batch.new()

  Build a batch from a stream:

      stream_of_args
      |> Stream.map(&MyWorker.new/1)
      |> Batch.new()

  Build a batch with callback options:

      Batch.new(list_of_jobs, batch_id: "abc-123", callback_worker: MyApp.CallbackWorker)
  """
  @spec new(Enumerable.t(changeset()), batch_opts()) :: t()
  def new(changesets, opts) when is_list_or_stream(changesets) and is_list(opts) do
    validate!(opts)

    opts =
      opts
      |> Keyword.put(:batch, true)
      |> Keyword.put_new(:callback_opts, [])
      |> Keyword.update!(:callback_opts, &Map.new/1)
      |> Keyword.put_new_lazy(:batch_id, &Oban.Pro.UUIDv7.generate/0)
      |> Map.new()

    changesets = Stream.map(changesets, &put_meta(&1, opts))

    %__MODULE__{changesets: changesets, opts: opts}
  end

  defp put_meta(changeset, opts) do
    meta =
      changeset
      |> Changeset.get_change(:meta, %{})
      |> Map.merge(opts)

    Changeset.put_change(changeset, :meta, meta)
  end

  defp validate!(opts) do
    Validation.validate_schema!(opts,
      batch_id: :string,
      batch_name: :string,
      callback_opts: :list,
      callback_worker: :module
    )
  end

  @doc """
  Add one or more jobs to a batch.

  ## Examples

  Add jobs to an existing batche:

      Batch.add(batch, Enum.map(args, &MyWorker.new/1))

  Add jobs to a batch one at a time:

      Batch.new()
      |> Batch.add(MyWorker.new(%{}))
      |> Batch.add(MyWorker.new(%{}))
      |> Batch.add(MyWorker.new(%{}))
  """
  @spec add(t(), Enumerable.t(changeset())) :: t()
  def add(batch, %Changeset{} = changeset), do: add(batch, [changeset])

  def add(%__MODULE__{} = batch, changesets) when is_list_or_stream(changesets) do
    changesets = Stream.map(changesets, &put_meta(&1, batch.opts))

    %{batch | changesets: Stream.concat(batch.changesets, changesets)}
  end

  @doc """
  Append to a batch from an existing batch job.

  The `batch_id` and any other batch options are propagated to the newly created batch.

  > #### Appending and Callbacks {: .warning}
  >
  > Batch callback jobs are _only inserted once_. Appending to a batch where callbacks were already
  > triggered, e.g. `batch_completed`, won't re-trigger the callback.

  ## Examples

  Build an empty batch from an existing batch job:

      job
      |> Batch.append()
      |> Batch.add(MyWorker.new(%{}))
      |> Batch.add(MyWorker.new(%{}))
      |> Oban.insert_all()

  Build a batch from an existing job with overriden options:

      Batch.append(job, callback_worker: MyApp.OtherWorker)
  """
  @spec append(Job.t(), batch_opts()) :: t()
  def append(%Job{meta: %{"batch_id" => _} = meta}, opts \\ []) when is_list(opts) do
    orig_opts =
      for {key, val} <- meta,
          key in ~w(batch_id batch_name callback_opts callback_worker),
          do: {String.to_existing_atom(key), val}

    orig_opts
    |> Keyword.merge(opts)
    |> new()
  end

  @doc """
  Get all non-callback jobs from a batch.

  ## Examples

  Fetch results from all jobs from within a `batch_completed/1` callback:

      def batch_completed(%Job{} = job) do
        results =
          job
          |> Batch.all_jobs()
          |> Enum.map(&fetch_recorded/1)

        ...
      end

  Get all batch jobs from anywhere using the `batch_id`:

      Batch.all_jobs("some-uuid-1234-5678")

  Get all jobs from anywhere with a custom Oban instance name:

      Batch.all_jobs(MyApp.Oban, "some-uuid-1234-5678")
  """
  @spec all_jobs(Oban.name(), Job.t() | String.t(), [repo_opts()]) :: [Job.t()]
  def all_jobs(name \\ Oban, job_or_batch_id, opts \\ [])

  def all_jobs(_name, %Job{conf: conf, meta: %{"batch_id" => batch_id}}, opts) do
    Repo.all(conf, batch_query(batch_id), opts)
  end

  def all_jobs(name, batch_id, opts) when is_binary(batch_id) do
    name
    |> Oban.config()
    |> Repo.all(batch_query(batch_id), opts)
  end

  def all_jobs(%Job{} = job, opts, []), do: all_jobs(Oban, job, opts)

  @doc """
  Cancel all non-callback jobs in a batch.

  Cancelling jobs via this function bulk-updates their state in the database without loading them
  into memory or resolving their workers. As a result, **worker `after_process/3` hooks are not
  called** for cancelled jobs. If you need to execute cleanup logic or other side effects when jobs
  are cancelled, consider using batch callbacks like `batch_cancelled/1` or implement cleanup logic
  at the application level.

  ## Examples

  Cancel jobs with the `job` in a `process/1` function:

      def process(job) do
        if should_stop_processing?(job.args) do
          Batch.cancel_jobs(job)
        else
          ...
        end
      end

  Cancel jobs from anywhere using the `batch_id`:

      Batch.cancel_jobs("some-uuid-1234-5678")

  Cancel jobs from anywhere with a custom Oban instance name:

      Batch.cancel_jobs(MyApp.Oban, "some-uuid-1234-5678")
  """
  @spec cancel_jobs(Oban.name(), Job.t() | String.t()) :: {:ok, non_neg_integer()}
  def cancel_jobs(oban_name \\ Oban, job_or_batch_id)

  def cancel_jobs(_oban_name, %Job{conf: conf, meta: %{"batch_id" => batch_id}}) do
    cancel_jobs(conf.name, batch_id)
  end

  def cancel_jobs(oban_name, batch_id) when is_binary(batch_id) do
    Oban.cancel_all_jobs(oban_name, batch_query(batch_id))
  end

  @doc """
  Create a batch from a workflow.

  All jobs in the workflow are augmented to also be part of a batch.

  ## Examples

  Create a batch from a workflow without any extra options:

      Workflow.new()
      |> Workflow.add(:step_1, MyWorker.new(%{id: 123}))
      |> Workflow.add(:step_2, MyWorker.new(%{id: 345}), deps: [:step_1])
      |> Workflow.add(:step_3, MyWorker.new(%{id: 456}), deps: [:step_2])
      |> Batch.from_workflow()
      |> Oban.insert_all()

  Create a batch with callback options:

      Batch.from_workflow(workflow, callback_opts: [priority: 3], callback_worker: MyApp.Worker)
  """
  @spec from_workflow(Workflow.t(), batch_opts()) :: t()
  def from_workflow(%Workflow{changesets: changesets}, opts \\ []) do
    new(changesets, opts)
  end

  @doc """
  Stream all non-callback jobs from a batch.

  ## Examples

  Stream all batch jobs from within a `batch_completed/1` callback:

      def batch_completed(%Job{} = job) do
        {:ok, account_ids} =
          MyApp.Repo.transaction(fn ->
            job
            |> Batch.stream_jobs()
            |> Enum.map(& &1.args["account_id"])
          end)

          # use_account_ids
      end

  Stream all batch jobs from anywhere using the `batch_id`:

      MyApp.Repo.transaction(fn ->
        "some-uuid-1234-5678"
        |> Batch.stream_jobs()
        |> Enum.map(& &1.args["account_id"])
      end)

  Stream all batch jobs using a custom Oban instance name:

      MyApp.Repo.transaction(fn ->
        MyApp.Oban
        |> Batch.stream_jobs("some-uuid-1234-5678")
        |> Enum.map(& &1.args["account_id"])
      end)
  """
  @spec stream_jobs(Oban.name(), Job.t() | String.t(), [repo_opts()]) :: Enum.t()
  def stream_jobs(name \\ Oban, job_or_batch_id, opts \\ [])

  def stream_jobs(_name, %Job{conf: conf, meta: %{"batch_id" => batch_id}}, opts) do
    Repo.stream(conf, batch_query(batch_id), opts)
  end

  def stream_jobs(name, batch_id, opts) when is_binary(batch_id) do
    name
    |> Oban.config()
    |> Repo.stream(batch_query(batch_id), opts)
  end

  def stream_jobs(%Job{} = job, opts, []), do: stream_jobs(Oban, job, opts)

  defp batch_query(batch_id) do
    Job
    |> where([j], fragment("? \\? 'batch_id'", j.meta))
    |> where([j], fragment("? ->> 'batch_id'", j.meta) == ^batch_id)
    |> where([j], is_nil(fragment("? ->> 'callback'", j.meta)))
    |> where([j], j.state in @all_states)
    |> order_by(asc: :id)
  end

  # Event Handling

  @doc false
  @impl Oban.Pro.Flusher
  def to_flush_mfa(job, conf) do
    case job.meta do
      %{"batch_id" => _, "callback" => _} -> :ignore
      %{"batch_id" => batch_id} -> {__MODULE__, :on_flush, [job, batch_id, conf]}
      _ -> :ignore
    end
  end

  @doc false
  def on_flush(job, batch_id, conf) do
    batch_worker = job.meta["batch_callback_worker"] || job.meta["callback_worker"] || job.worker

    with {:ok, worker} <- Worker.from_string(batch_worker),
         supported = supported_callbacks(worker),
         {:ok, {states, exists}} <- states_for_callbacks(supported, batch_id, conf) do
      for callback <- supported,
          callback not in exists,
          callback_ready?(callback, states) do
        insert_callback(callback, worker, job, conf)
      end
    end
  end

  defp supported_callbacks(worker) do
    for {name, [new, old]} <- @callbacks_to_combined,
        function_exported?(worker, new, 1) or function_exported?(worker, old, 1),
        do: name
  end

  defp states_for_callbacks([], _batch_id, _conf), do: :ok

  defp states_for_callbacks(callbacks, batch_id, conf) do
    state_query =
      callbacks
      |> Enum.flat_map(&Map.fetch!(@callbacks_to_states, &1))
      |> Enum.uniq()
      |> Enum.reduce(:none, fn state, acc ->
        query =
          Job
          |> select([_], [type(^state, :string)])
          |> where([j], fragment("? \\? 'batch_id'", j.meta))
          |> where([j], fragment("? ->> 'batch_id'", j.meta) == ^batch_id)
          |> where([j], is_nil(fragment("? ->> 'callback'", j.meta)))
          |> where([j], j.state == ^state)
          |> limit(1)

        if acc == :none, do: query, else: union_all(acc, ^query)
      end)

    exist_query =
      Enum.reduce(callbacks, :none, fn callback, acc ->
        query =
          Job
          |> select([_j], [type(^callback, :string)])
          |> where([j], fragment("? \\? 'batch_id'", j.meta))
          |> where([j], fragment("? ->> 'batch_id'", j.meta) == ^batch_id)
          |> where([j], fragment("? ->> 'callback'", j.meta) == ^callback)
          |> where([j], j.state in @all_states)

        if acc == :none, do: query, else: union_all(acc, ^query)
      end)

    Repo.transaction(conf, fn ->
      states = conf |> Repo.all(state_query) |> List.flatten()
      exists = conf |> Repo.all(exist_query) |> List.flatten()

      {states, exists}
    end)
  end

  defp callback_ready?(callback, batch_states) do
    ready? =
      @callbacks_to_states
      |> Map.fetch!(callback)
      |> Enum.any?(&(&1 in batch_states))

    # Other callbacks use a negated query to avoid counting `completed` jobs.
    if callback in ~w(cancelled discarded retryable) do
      ready?
    else
      not ready?
    end
  end

  defp insert_callback(callback, worker, job, conf) do
    call_opts =
      job.meta
      |> Map.get("callback_opts", %{})
      |> Map.put_new("args", Map.get(job.meta, "batch_callback_args", %{}))
      |> Map.put_new("meta", Map.get(job.meta, "batch_callback_meta", %{}))
      |> Map.put_new("queue", Map.get(job.meta, "batch_callback_queue", job.queue))

    unique = [
      period: :infinity,
      fields: [:worker, :queue, :meta],
      keys: [:batch_id, :callback],
      states: Job.states()
    ]

    {args, call_opts} = Map.pop(call_opts, "args")
    {meta, call_opts} = Map.pop(call_opts, "meta")

    xtra_meta =
      job.meta
      |> Map.take(~w(batch_id batch_name))
      |> Map.put("callback", callback)
      |> Map.merge(meta)

    opts =
      call_opts
      |> Keyword.new(fn {key, val} -> {String.to_existing_atom(key), val} end)
      |> Keyword.put(:meta, xtra_meta)
      |> Keyword.put(:unique, unique)

    changeset = worker.new(args, opts)

    if not changeset.valid? and Keyword.has_key?(changeset.errors, :args) do
      changeset
      |> structured_error_message()
      |> Logger.error()
    else
      {:ok, Smart.insert_job(conf, changeset, [])}
    end
  end

  defp structured_error_message(changeset) do
    """
    [Oban.Pro.Batch] can't insert batch callback because it has invalid keys:

      #{get_in(changeset.errors, [:args, Access.elem(0)])}

    Use one of the following options to restore batch callbacks:

    * Modify structured `keys` or `required` to allow the missing keys
    * Include the required arguments with the `batch_callback_args` option
    * Specify a different callback worker with the `batch_callback_worker` option
    """
  end
end