kvasir_agent_server

defmodule Kvasir.AgentServer do
  @moduledoc ~S"""
  Documentation for `Kvasir.AgentServer`.
  """

  @type id :: atom | reference()
  @type status :: :unknown | :ready | :reconfiguring | :starting | :stopping
  @type agent :: %{
          server: id,
          id: String.t(),
          agent: module,
          partition: String.t(),
          opts: Keyword.t()
        }

  @default_control_port 9393
  @default_agent_ports [
    9397,
    9399,
    9733,
    9737,
    9739,
    9773,
    9777,
    9779,
    9793,
    9797,
    9799,
    9933,
    9937,
    9939,
    9973,
    9977,
    9979,
    9993,
    9997,
    9999
  ]

  @doc """
  Default control port for Kvasir AgentServer.

  Commonly: `#{@default_control_port}`
  """
  @spec default_control_port :: pos_integer()
  def default_control_port, do: @default_control_port

  @doc """
  Default ports to use for agents command handlers.

  Commonly: #{@default_agent_ports |> Enum.map(&"\n- `#{&1}`") |> Enum.join()}
  """
  @spec default_agent_ports :: [pos_integer()]
  def default_agent_ports, do: @default_agent_ports

  use Supervisor
  alias Kvasir.AgentServer.Config
  alias Kvasir.AgentServer.SystemMonitor
  @opts ~w(warmup monitor_system custom_commands advertise_host)a

  @doc @moduledoc
  @spec child_spec(opts :: Keyword.t()) :: Supervisor.child_spec()
  def child_spec(opts \\ []) do
    server = opts[:id] || make_ref()
    agents = agents(opts[:agents])
    o = Keyword.take(opts, @opts)

    %{
      id: server,
      restart: :permanent,
      shutdown: :infinity,
      type: :supervisor,
      modules: [__MODULE__],
      start: {__MODULE__, :start_link, [server, agents, o]}
    }
  end

  @doc false
  @spec start_link(id :: id, agents :: [agent], Keyword.t()) :: Supervisor.on_start()
  def start_link(id, agents, opts \\ [])

  def start_link(id, agents, opts) do
    ensure_ranch!()
    monitor? = opts[:monitor_system] || false

    if monitor?, do: SystemMonitor.start()

    ip =
      if host = opts[:advertise_host] do
        host
      else
        {:ok, ips} = :inet.getif()
        {ip, _, _} = List.first(ips)
        ip
      end

    o = opts |> Keyword.take(@opts) |> Keyword.update(:custom_commands, %{}, &custom_commands!/1)

    agents =
      agents
      |> Enum.zip(default_agent_ports())
      |> Enum.map(fn {agent, p} ->
        if monitor? do
          if dir = Keyword.get(agent.agent.config(:cache, []), :directory) do
            SystemMonitor.register(dir, agent)
          end
        end

        agent
        |> Map.put(:server, id)
        |> Map.update!(:opts, &(&1 |> Keyword.put_new(:ip, ip) |> Keyword.put_new(:port, p)))
      end)

    if is_atom(id) do
      Supervisor.start_link(__MODULE__, {id, agents, o}, name: id)
    else
      Supervisor.start_link(__MODULE__, {id, agents, o})
    end
  end

  @impl Supervisor
  def init({id, agents, opts}) do
    :ets.new(id, [:named_table, :set, :public, read_concurrency: true])
    Config.set_status(id, :reconfiguring)
    Config.set_agents(id, agents)

    children = [
      {__MODULE__.Config, server: id},
      {__MODULE__.Control.Handler, server: id, custom_commands: opts[:custom_commands]},
      {__MODULE__.AgentManager,
       server: id, agents: agents, warmup: Keyword.get(opts, :warmup, false)},
      %{id: :shutdown_monitor, type: :worker, start: {__MODULE__, :monitor, [id]}}
    ]

    Supervisor.init(children, strategy: :rest_for_one)
  end

  def monitor(id) do
    {:ok,
     spawn_link(fn ->
       Process.flag(:trap_exit, true)
       do_monitor(id)
     end)}
  end

  defp do_monitor(id) do
    receive do
      {:EXIT, _, reason} ->
        Config.set_status(id, :stopping)
        require Logger
        Logger.info(fn -> "Kvasir AgentServer<#{inspect(id)}>: Stopping #{inspect(reason)}." end)

      _ ->
        do_monitor(id)
    end
  end

  @spec ensure_ranch! :: :ok | no_return()
  defp ensure_ranch! do
    case :application.start(:ranch) do
      :ok -> :ok
      {:error, {:already_started, :ranch}} -> :ok
    end
  end

  @spec agents(term) :: [agent] | no_return()
  defp agents(nil), do: raise("Missing `:agents` in child spec options.")

  defp agents(agent) when is_map(agent) do
    agent |> Enum.to_list() |> agents()
  end

  defp agents(agents), do: Enum.map(agents, &agent_entry/1)

  defp agent_entry({id, agent}),
    do: %{
      id: id,
      agent: agent,
      partition: "*",
      counter: :counters.new(1, [:write_concurrency]),
      opts: []
    }

  defp agent_entry({id, agent, opts}) do
    {partition, o} = Keyword.pop(opts, :partition)

    p = partition || "*"

    %{
      id: id,
      agent: agent,
      partition: p,
      counter: :counters.new(1, [:write_concurrency]),
      opts: o
    }
  end

  defp agent_entry(agent = %{}) do
    id = agent[:id] || agent["id"] || raise "Missing agent id: #{inspect(agent)}"

    agent =
      case agent[:agent] || agent["agent"] do
        a when is_atom(a) -> a
        a when is_binary(a) -> String.to_existing_atom(a)
        _ -> raise "Missing agent: #{inspect(agent)}"
      end

    partition = agent[:partition] || agent["partition"] || "*"

    opts =
      Enum.map(agent[:opts] || agent["opts"] || [], fn
        {k, v} when is_atom(k) -> {k, v}
        {k, v} -> {String.to_existing_atom(k), v}
      end)

    %{
      id: id,
      agent: agent,
      partition: partition,
      counter: :counters.new(1, [:write_concurrency]),
      opts: opts
    }
  end

  alias Kvasir.AgentServer.Control.Server, as: ControlServer
  defp custom_commands!(opts)

  defp custom_commands!(nil), do: %{}

  defp custom_commands!(commands) do
    commands
    |> Enum.map(fn {command, call} ->
      cmd = command |> clean_command() |> valid_command!()

      unless is_function(call, 1) do
        raise "Kvasir AgentsServer: Invalid custom command \"#{cmd}\" has not proper call."
      end

      {cmd, call}
    end)
    |> Map.new()
  end

  defp clean_command(command), do: command |> to_string() |> String.trim() |> String.upcase()

  defp valid_command!(command) do
    if command in ControlServer.commands() or not Regex.match?(~r/^[A-Z\_0-9]+$/, command) do
      raise "Kvasir AgentsServer: Invalid custom command \"#{command}\"."
    end

    command
  end
end