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

  @behaviour Oban.Pro.Stage

  import Ecto.Query, only: [where: 3]

  alias Oban.{Job, Repo}
  alias Oban.Pro.Utils

  require Logger

  @hooks [
    after_cancelled: 2,
    before_new: 2,
    before_process: 1,
    after_process: 3,
    on_cancelled: 2,
    on_discarded: 2
  ]

  defstruct modules: []

  @impl Oban.Pro.Stage
  def init(worker, modules) do
    {:ok, struct!(__MODULE__, modules: [worker | modules])}
  end

  @impl Oban.Pro.Stage
  def before_new(args, opts, %{modules: modules}) do
    (modules ++ get_hooks())
    |> Enum.filter(&hook_exported?(&1, :before_new))
    |> Enum.reduce_while({:ok, args, opts}, fn module, {:ok, args, opts} ->
      case module.before_new(args, opts) do
        {:ok, _, _} = ok -> {:cont, ok}
        other -> {:halt, other}
      end
    end)
  end

  @impl Oban.Pro.Stage
  def before_process(job, %{modules: modules}) do
    (modules ++ get_hooks())
    |> Enum.filter(&hook_exported?(&1, :before_process))
    |> Enum.reduce_while({:ok, job}, fn module, {:ok, job} ->
      case module.before_process(job) do
        {:ok, _} = ok -> {:cont, ok}
        other -> {:halt, other}
      end
    end)
  end

  def attach_hook(module) do
    with :ok <- Utils.validate_opts([module], &validate_module/1) do
      modules =
        get_hooks()
        |> Enum.concat([module])
        |> Enum.uniq()

      :persistent_term.put(:oban_pro_hooks, modules)
    end
  end

  def detach_hook(module) do
    case get_hooks() do
      [_ | _] = modules ->
        :persistent_term.put(:oban_pro_hooks, modules -- [module])

      _ ->
        :ok
    end
  end

  # The `Workflow.after_cancelled` hook is deprecated and will be removed in a future version.
  def execute_after_cancelled(reason, job) do
    with {worker, job, opts} <- resolve_processing(job) do
      execute_hook(worker, opts, :on_cancelled, [:dependency, job])
      execute_hook(worker, opts, :after_cancelled, [reason, job])
    end
  end

  def execute_on_cancelled(reason, job) do
    with {worker, job, opts} <- resolve_processing(job) do
      execute_hook(worker, opts, :on_cancelled, [reason, job])
    end
  end

  def execute_on_discarded(reason, job) do
    with {worker, job, opts} <- resolve_processing(job) do
      execute_hook(worker, opts, :on_discarded, [reason, job])
    end
  end

  def handle_event([:oban, :job, _], _timing, meta, _conf) do
    with {worker, job, opts} <- resolve_processing(meta.job) do
      execute_hook(worker, opts, :after_process, [to_hook_state(meta.state), job, meta[:result]])
    end
  end

  def handle_event([:oban, :engine, :cancel_all_jobs, _], _timing, meta, _conf) do
    global? = Enum.any?(get_hooks(), &hook_exported?(&1, :on_cancelled))

    Task.start(fn ->
      load_and_exec = fn partials ->
        ids = Enum.map(partials, & &1.id)

        meta.conf
        |> Repo.all(where(Job, [j], j.id in ^ids))
        |> Enum.each(&execute_on_cancelled(:manual, &1))
      end

      meta.jobs
      |> Stream.filter(&(global? or hook_exported?(&1, :on_cancelled)))
      |> Stream.chunk_every(100)
      |> Stream.each(load_and_exec)
      |> Stream.run()
    end)
  end

  defp execute_hook(worker, opts, hook, args) do
    mod_hooks = Keyword.get(opts, :hooks, [])
    all_hooks = Enum.uniq([worker | mod_hooks] ++ get_hooks())

    for module <- all_hooks, hook_exported?(module, hook) do
      apply(module, hook, args)
    end
  catch
    kind, value ->
      Logger.error(fn ->
        "[Oban.Pro.Worker] hook error: " <> Exception.format(kind, value, __STACKTRACE__)
      end)
  end

  defp resolve_processing(job) do
    with nil <- Process.get(:oban_processing),
         {:ok, worker} <- Oban.Worker.from_string(job.worker),
         opts = worker.__opts__(),
         {:ok, job} <- Oban.Pro.Worker.before_hook(worker, job, opts) do
      {worker, job, opts}
    end
  end

  defp get_hooks, do: :persistent_term.get(:oban_pro_hooks, [])

  defp to_hook_state(:cancelled), do: :cancel
  defp to_hook_state(:success), do: :complete
  defp to_hook_state(:discard), do: :discard
  defp to_hook_state(:failure), do: :error
  defp to_hook_state(:snoozed), do: :snooze

  defp hook_exported?(worker, fun) when is_binary(worker) do
    case Oban.Worker.from_string(worker) do
      {:ok, module} -> hook_exported?(module, fun)
      _ -> false
    end
  end

  defp hook_exported?(module, fun) when is_atom(module) do
    module
    |> cached_hooks()
    |> MapSet.member?(fun)
  end

  defp hook_exported?(%{worker: worker}, fun), do: hook_exported?(worker, fun)
  defp hook_exported?(_map, _fun), do: false

  defp cached_hooks(module) do
    case :persistent_term.get({:oban_pro_hook_cache, module}, :not_cached) do
      :not_cached ->
        module
        |> discover_hooks()
        |> tap(&:persistent_term.put({:oban_pro_hook_cache, module}, &1))

      hooks ->
        hooks
    end
  end

  defp discover_hooks(module) do
    if Code.ensure_loaded?(module) do
      for {fun, arity} <- @hooks, function_exported?(module, fun, arity), into: MapSet.new() do
        fun
      end
    else
      MapSet.new()
    end
  end

  # Validation

  defp validate_module(module) do
    cond do
      not Code.ensure_loaded?(module) ->
        {:error, "unable to load callback module #{inspect(module)}"}

      Enum.any?(@hooks, &function_exported?(module, elem(&1, 0), elem(&1, 1))) ->
        :ok

      true ->
        names = Enum.map_join(@hooks, ", ", fn {fun, arity} -> "#{fun}/#{arity}" end)

        {:error, "#{inspect(module)} doesn't implement known hooks: #{names}"}
    end
  end
end
