defmodule Oban.Pro.Diagnostics do
  @moduledoc false

  use GenServer

  alias __MODULE__, as: State
  alias Oban.Notifier

  defstruct [:timer, interval: :timer.seconds(5), subscribed: MapSet.new()]

  @info_keys ~w(
      current_stacktrace
      heap_size
      memory
      message_queue_len
      reductions
      stack_size
      status
    )a

  @doc false
  def start_link(opts \\ []) do
    state = struct!(State, opts)

    GenServer.start_link(__MODULE__, state, name: __MODULE__)
  end

  # Callbacks

  @impl GenServer
  def init(state) do
    Process.flag(:trap_exit, true)

    {:ok, schedule_discover(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_call(:discover, _from, state) do
    {:reply, :ok, discover_and_subscribe(state)}
  end

  @impl GenServer
  def handle_info(:discover, state) do
    {:noreply, discover_and_subscribe(state)}
  end

  def handle_info({:notification, :diagnostics, %{"job_id" => job_id}}, state) do
    finder = fn conf ->
      case Registry.lookup(Oban.Registry, {conf.name, {:executing, job_id}}) do
        [{pid, _}] -> {conf, pid}
        _ -> nil
      end
    end

    case Enum.find_value(list_oban_instances(), finder) do
      {conf, pid} ->
        payload = %{"info" => process_info(pid), "job_id" => job_id, "node" => inspect(conf.node)}

        Notifier.notify(conf, :diagnostics_reply, payload)

      nil ->
        :ok
    end

    {:noreply, state}
  end

  def handle_info(_message, state) do
    {:noreply, state}
  end

  # Helpers

  defp discover_and_subscribe(state) do
    state =
      Enum.reduce(list_oban_instances(), state, fn conf, acc ->
        if MapSet.member?(acc.subscribed, conf.name) do
          acc
        else
          :ok = Notifier.listen(conf.name, [:diagnostics])

          %{acc | subscribed: MapSet.put(acc.subscribed, conf.name)}
        end
      end)

    state
    |> cancel_timer()
    |> schedule_discover()
  end

  defp list_oban_instances do
    match = [{{{:"$1", Oban.Notifier}, :_, :_}, [], [:"$1"]}]

    Oban.Registry
    |> Registry.select(match)
    |> Enum.map(&Oban.config/1)
  end

  defp process_info(pid) do
    case Process.info(pid, @info_keys) do
      nil ->
        %{}

      info ->
        info
        |> Keyword.update!(:current_stacktrace, &format_stacktrace/1)
        |> Keyword.update!(:status, &to_string/1)
        |> Map.new(fn {key, val} -> {to_string(key), val} end)
    end
  end

  defp format_stacktrace(nil), do: ""
  defp format_stacktrace(stacktrace), do: Exception.format_stacktrace(stacktrace)

  defp cancel_timer(%{timer: timer} = state) when is_reference(timer) do
    Process.cancel_timer(timer)

    %{state | timer: nil}
  end

  defp cancel_timer(state), do: state

  defp schedule_discover(state) do
    timer = Process.send_after(self(), :discover, state.interval)

    %{state | timer: timer}
  end
end