defmodule Elsa.Group.Manager.WorkerSupervisor do
  @moduledoc """
  Provides functions to encapsulate the management of worker
  processes by the consumer group manager.
  """
  use Supervisor

  import Record, only: [defrecord: 2, extract: 2]
  import Elsa.ElsaSupervisor, only: [registry: 1]

  alias Elsa.Consumer.Worker
  alias Elsa.ElsaRegistry
  alias Elsa.Group.Acknowledger
  alias Elsa.Group.Manager
  alias Elsa.Group.Manager.State

  alias Elsa.Util

  require Logger

  defrecord :brod_received_assignment, extract(:brod_received_assignment, from_lib: "brod/include/brod.hrl")

  # ID used for the underlying DynamicSupervisor.
  # Defined as a module attribute to avoid typo bugs.
  @dynamic_supervisor_id :worker_dynamic_supervisor

  defmodule WorkerState do
    @moduledoc """
    Tracks the running state of the worker process from the perspective of the group manager.
    """
    defstruct [:pid, :ref, :generation_id, :topic, :partition, :latest_offset]
  end

  @doc """
  Start the main Supervisor.  On init this will start a single DynamicSupervisor as a child.

  Options:
    :connection - Elsa.connection that will identify this supervisor in the Elsa registry.
  """
  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts) do
    connection = Keyword.fetch!(opts, :connection)
    Supervisor.start_link(__MODULE__, opts, name: {:via, ElsaRegistry, {registry(connection), __MODULE__}})
  end

  @impl true
  @spec init(keyword()) :: {:ok, {Supervisor.sup_flags(), [Supervisor.child_spec()]}}
  def init(opts) do
    children = [
      %{
        id: @dynamic_supervisor_id,
        start: {DynamicSupervisor, :start_link, [[]]},
        restart: :transient
      }
    ]

    Supervisor.init(children,
      strategy: :one_for_one,
      # Normal restarts are pretty common, for example, if the group coordinator isn't available yet.
      # So we need to allow substantially more than the default 3 restart in 5 seconds.
      max_restarts: Keyword.get(opts, :worker_supervisor_max_restarts, 30),
      max_seconds: Keyword.get(opts, :worker_supervisor_max_seconds, 5)
    )
  end

  @doc """
  Retrieve the generation id, used in tracking assignments of workers to topic/partition,
  from the worker state map.
  """
  @spec get_generation_id(map(), Elsa.topic(), Elsa.partition()) :: Manager.generation_id()
  def get_generation_id(workers, topic, partition) do
    Map.get(workers, {topic, partition})
    |> Map.get(:generation_id)
  end

  @doc """
  Terminate all workers under the supervisor, giving them a chance to exit gracefully.

  Options:
    :restart_supervisor - if true, this function will restart the dynamic supervisor after killing it off.
      Default true.
  """
  @spec stop_all_workers(Elsa.connection(), map(), list()) :: map()
  def stop_all_workers(connection, workers, options \\ []) do
    Logger.info("Stopping all workers for #{inspect(connection)}")

    # Demonitor all the workers, so that our client doesn't get a bunch of :EXIT signals
    workers
    |> Map.values()
    |> Enum.each(fn worker ->
      Logger.info("#{__MODULE__}: Gracefully stopping group worker #{inspect(worker)}")
      Process.demonitor(worker.ref)
    end)

    module_supervisor = get_worker_supervisor(connection)

    dynamic_worker_supervisor = get_dynamic_supervisor(module_supervisor)

    if dynamic_worker_supervisor != :undefined do
      # Synchronously stops the DynamicSupervisor and its children
      DynamicSupervisor.stop(dynamic_worker_supervisor)
    else
      Logger.warning(
        "#{__MODULE__}: Attempted to stop #{@dynamic_supervisor_id}, but it does not exist under Supervisor #{inspect(module_supervisor)}"
      )
    end

    # Make sure the DynamicSupervisor itself is truly cleaned up from the Supervisor's perspective,
    # so that it will restart reliably
    :ok = Supervisor.terminate_child(module_supervisor, @dynamic_supervisor_id)

    # Restart the dynamic supervisor
    if Keyword.get(options, :restart_supervisor, true) do
      {:ok, _pid} = Supervisor.restart_child(module_supervisor, @dynamic_supervisor_id)
      :ok
    end

    %{}
  end

  @doc """
  Restart the specified worker from the manager state. Retrieve the latest recorded
  offset and pass it to the new worker to pick up where the previous left off if it
  has been recorded.
  """
  @spec restart_worker(map(), reference(), struct()) :: map()
  def restart_worker(workers, ref, %State{} = state) do
    worker = get_by_ref(workers, ref)

    latest_offset =
      Acknowledger.get_latest_offset(
        {:via, ElsaRegistry, {registry(state.connection), Acknowledger}},
        worker.topic,
        worker.partition
      ) || worker.latest_offset

    assignment = brod_received_assignment(topic: worker.topic, partition: worker.partition, begin_offset: latest_offset)

    start_worker(workers, worker.generation_id, assignment, state)
  end

  @doc """
  Construct an argument payload for instantiating a worker process, generate a
  topic/partition assignment and instantiate the worker process with both under
  the dynamic supervisor. Record the manager-relevant information and store in the
  manager state map tracking active worker processes.
  """
  @spec start_worker(map(), integer(), tuple(), struct()) :: map()
  def start_worker(workers, generation_id, assignment, %State{} = state) do
    assignment = Enum.into(brod_received_assignment(assignment), %{})

    init_args = [
      topic: assignment.topic,
      partition: assignment.partition,
      generation_id: generation_id,
      begin_offset: assignment.begin_offset,
      handler: state.handler,
      handler_init_args: state.handler_init_args,
      connection: state.connection,
      config: state.config
    ]

    Logger.info("#{__MODULE__}: Starting group consumer worker: #{inspect(init_args)}")

    # If we have an assigned offset, go ahead and update our bookkeeping in the Acknowledger.
    # This way if the worker needs to be restarted before acking anything,
    # we won't rewind to some stale offset.
    if assignment.begin_offset != :undefined do
      Acknowledger.set_latest_offset(
        {:via, ElsaRegistry, {registry(state.connection), Acknowledger}},
        assignment.topic,
        assignment.partition,
        assignment.begin_offset
      )
    end

    dynamic_supervisor = get_dynamic_supervisor(get_worker_supervisor(state.connection))

    {:ok, worker_pid} = DynamicSupervisor.start_child(dynamic_supervisor, {Worker, init_args})
    ref = Process.monitor(worker_pid)

    new_worker = %WorkerState{
      pid: worker_pid,
      ref: ref,
      generation_id: generation_id,
      topic: assignment.topic,
      partition: assignment.partition,
      latest_offset: assignment.begin_offset
    }

    Logger.info("Group consumer worker started: #{inspect(new_worker)}")

    Map.put(workers, {assignment.topic, assignment.partition}, new_worker)
  end

  defp get_by_ref(workers, ref) do
    workers
    |> Map.values()
    |> Enum.find(fn worker -> worker.ref == ref end)
  end

  defp get_worker_supervisor(connection) do
    Util.with_registry(connection, fn registry ->
      ElsaRegistry.whereis_name({registry, __MODULE__})
    end)
  end

  defp get_dynamic_supervisor(worker_supervisor_pid) do
    [{@dynamic_supervisor_id, dynamic_worker_supervisor, _, _}] = Supervisor.which_children(worker_supervisor_pid)
    dynamic_worker_supervisor
  end
end