defmodule Oban.Pro.Stages.Chain do
  @moduledoc false

  @behaviour Oban.Pro.Flusher
  @behaviour Oban.Pro.Stage

  import Ecto.Query

  alias Ecto.Changeset
  alias Oban.{Job, Repo, Validation}
  alias Oban.Pro.Utils

  @hold_date ~U[3000-01-01 00:00:00.000000Z]

  defstruct [:by, :worker, on_cancelled: :ignore, on_discarded: :ignore]

  defmacrop drop_hold(meta) do
    quote do
      fragment(~s(? || '{"on_hold":false, "uniq_bmp":[]}'), unquote(meta))
    end
  end

  defmacrop continue_state(prefix, meta) do
    quote do
      fragment(
        """
        (CASE WHEN ? \\? 'orig_scheduled_at' THEN 'scheduled' ELSE 'available' END)::?.oban_job_state
        """,
        unquote(meta),
        unquote(prefix)
      )
    end
  end

  defmacrop continue_at(meta, now) do
    quote do
      fragment(
        """
        CASE
        WHEN ? \\? 'orig_scheduled_at' THEN timezone('utc', to_timestamp(div((?->'orig_scheduled_at')::bigint, 1000000)::float))
        ELSE ?
        END
        """,
        unquote(meta),
        unquote(meta),
        unquote(now)
      )
    end
  end

  defmacrop same_chain(meta_a, meta_b) do
    quote do
      fragment("?->>'chain_id' = ?->>'chain_id'", unquote(meta_a), unquote(meta_b))
    end
  end

  @spec chain?(Job.changeset()) :: boolean()
  def chain?(changeset), do: match?(%{changes: %{meta: %{"chain_id" => _}}}, changeset)

  @spec get_key(Job.changeset()) :: nil | String.t()
  def get_key(%{changes: %{meta: %{"chain_id" => chain_id}}}), do: chain_id
  def get_key(_changeset), do: nil

  @spec continuable?(String.t(), map()) :: boolean()
  def continuable?(state, meta) do
    state == "completed" or
      (state == "cancelled" and meta["on_cancelled"] != "hold") or
      (state == "discarded" and meta["on_discarded"] != "hold")
  end

  @spec to_postponed(Job.changeset()) :: Job.changeset()
  def to_postponed(changeset) do
    meta =
      if Changeset.get_field(changeset, :state) == "scheduled" do
        orig_at =
          changeset
          |> Changeset.get_change(:scheduled_at)
          |> DateTime.to_unix(:microsecond)

        %{"on_hold" => true, "orig_scheduled_at" => orig_at}
      else
        %{"on_hold" => true}
      end

    changeset
    |> Changeset.update_change(:meta, &Map.merge(&1, meta))
    |> Changeset.put_change(:state, "scheduled")
    |> Changeset.put_change(:scheduled_at, @hold_date)
  end

  @doc false
  def rescue_chains(conf, opts \\ []) do
    prev_query =
      Job
      |> where([j], j.state in ~w(available scheduled retryable executing))
      |> where([j], fragment("? \\? 'chain_id'", j.meta))
      |> where([j], same_chain(j.meta, parent_as(:jobs).meta))
      |> where([j], fragment("?->>'on_hold'", j.meta) in ~w(true false))
      |> where([j], j.id < parent_as(:jobs).id)

    subquery =
      from(j in Job, as: :jobs)
      |> where([j], j.state == "scheduled")
      |> where([j], fragment("? \\? 'chain_id'", j.meta))
      |> where([j], fragment("?->>'on_hold' = 'true'", j.meta))
      |> where([j], not exists(prev_query))
      |> order_by(asc: :id)
      |> limit(1)

    query =
      Job
      |> join(:inner, [j], x in subquery(subquery), on: j.id == x.id)
      |> update([j],
        set: [
          meta: drop_hold(j.meta),
          state: continue_state(literal(^conf.prefix), j.meta),
          scheduled_at: continue_at(j.meta, ^DateTime.utc_now())
        ]
      )

    Repo.update_all(conf, query, [], opts)

    :ok
  end

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

    Utils.hash64(chain_data)
  end

  # Stage

  @impl Oban.Pro.Stage
  def init(worker, opts) do
    with :ok <- validate(opts) do
      opts =
        opts
        |> Keyword.put_new(:by, :worker)
        |> Keyword.update!(:by, &Utils.normalize_by/1)
        |> Keyword.put(:worker, worker)

      {:ok, struct!(__MODULE__, opts)}
    end
  end

  @impl Oban.Pro.Stage
  def before_new(args, opts, conf) do
    conf = merge_conf_opts(conf, opts)

    meta = %{
      "chain" => true,
      "chain_id" => to_chain_id(conf.by, conf.worker, args, opts),
      "on_cancelled" => to_string(conf.on_cancelled),
      "on_discarded" => to_string(conf.on_discarded),
      "on_hold" => false
    }

    opts = Keyword.update(opts, :meta, meta, &Map.merge(&1, meta))

    {:ok, args, opts}
  end

  # Flushing

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

  def on_flush(%{state: state, meta: %{"chain_id" => chain_id} = meta}, conf) do
    if continuable?(state, meta) do
      subquery =
        Job
        |> select([j], j.id)
        |> where([j], j.state == "scheduled")
        |> where([j], fragment("? \\? 'chain_id'", j.meta))
        |> where([j], fragment("?->>'chain_id'", j.meta) == ^chain_id)
        |> where([j], fragment("?->>'on_hold'", j.meta) == "true")
        |> order_by(asc: :id)
        |> limit(1)
        |> lock("FOR UPDATE")

      query =
        Job
        |> join(:inner, [j], x in subquery(subquery), on: j.id == x.id)
        |> update([j],
          set: [
            meta: drop_hold(j.meta),
            state: continue_state(literal(^conf.prefix), j.meta),
            scheduled_at: continue_at(j.meta, ^DateTime.utc_now())
          ]
        )

      Repo.update_all(conf, query, [])
    end
  end

  # Helpers

  defp merge_conf_opts(conf, opts) do
    case Keyword.fetch(opts, :chain) do
      {:ok, chain} ->
        Enum.reduce(chain, conf, fn {key, val}, acc -> Map.replace!(acc, key, val) end)

      :error ->
        conf
    end
  end

  defp take_keys(map, keys) do
    Enum.map(keys, fn key -> inspect(map[key] || map[to_string(key)]) end)
  end

  defp validate(opts) do
    Validation.validate_schema(opts,
      by: {:custom, &Oban.Pro.Validation.validate_by/1},
      on_cancelled: {:enum, [:hold, :ignore]},
      on_discarded: {:enum, [:hold, :ignore]}
    )
  end
end
