kvasir_agent_server

defmodule Kvasir.AgentServer.Control.Server do
  @commands ~W(REBUILD STATUS CONNECT METRICS VERIFY_CACHE)

  def commands, do: @commands

  def init(socket, transport, server) do
    # Join for status updates
    :pg.join({server, :status}, spawn_link(fn -> status_loop(socket, transport, server) end))

    :ok
  end

  def close(_server) do
    :ok
  end

  ### The Commands
  alias Kvasir.AgentServer.Config

  @spec handle_command([String.t()], term, %{required(String.t()) => fun}) ::
          :ok | {:reply, iodata()}
  def handle_command(command, server, _custom_commands)

  def handle_command(["REBUILD", topic | opts], server, _custom_commands) do
    [key | part] = :lists.reverse(opts)
    partition = part |> :lists.reverse() |> Enum.join(" ")

    if agent =
         Enum.find_value(
           Config.agents(server, topic),
           &if(&1.partition == partition, do: &1.agent)
         ) do
      with {:ok, p_key} <- agent.__agent__(:key).parse(key),
           true <- match_partition?(partition, p_key) || {:error, :key_partition_mismatch},
           :ok <- agent.rebuild(p_key) do
        {:reply, "OK\n"}
      else
        {:error, err} -> {:reply, "ERR #{inspect(err)}\n"}
      end
    else
      {:reply, "ERR unknown_agent\n"}
    end
  end

  def handle_command(["VERIFY_CACHE", restore, topic | opts], server, _custom_commands) do
    case String.downcase(restore) do
      r when r in ~W(true yes y 1) -> check_cache(topic, opts, server, true)
      r when r in ~W(false no n 0) -> check_cache(topic, opts, server, false)
      _ -> {:reply, "ERR invalid restore argument <true|false> required\n"}
    end
  end

  def handle_command(["CONNECT" | opts], server, _custom_commands) do
    if id = List.first(opts) do
      a =
        server
        |> Config.agents(id)
        |> Enum.map(
          &[id, ?\ , ip!(&1.opts[:ip]), ?\ , to_string(&1.opts[:port]), ?\ , &1.partition, ?\n]
        )

      {:reply,
       [
         ["LIST ", id, ?\n],
         a,
         ["DONE ", id, ?\n]
       ]}
    else
      a =
        server
        |> Config.agents()
        |> Enum.map(
          &[&1.id, ?\ , ip!(&1.opts[:ip]), ?\ , to_string(&1.opts[:port]), ?\ , &1.partition, ?\n]
        )

      {:reply,
       [
         ["LIST", ?\n],
         a,
         ["DONE", ?\n]
       ]}
    end
  end

  def handle_command(["METRICS"], server, _custom_commands) do
    a =
      server
      |> Config.agents()
      |> Enum.map(&[&1.id, ?\ , &1.partition, ?\ , to_string(:counters.get(&1.counter, 1)), ?\n])

    {:reply,
     [
       "METRICS\n",
       a,
       "DONE METRICS\n"
     ]}
  end

  def handle_command(["STATUS"], server, _custom_commands),
    do: {:reply, ["STATUS ", to_string(Config.status(server)), ?\n]}

  def handle_command([cmd | opts], server, custom_commands) do
    if c = custom_commands[cmd], do: c.(opts), else: unknown!([cmd | opts], server)
  end

  def handle_command(unknown, server, _custom_commands), do: unknown!(unknown, server)

  defp unknown!(unknown, server) do
    require Logger

    Logger.warn(fn ->
      "Kvasir AgentServer<#{inspect(server)}>: Received unknown control command: #{
        inspect(unknown)
      } "
    end)

    if List.first(unknown) in @commands do
      {:reply, "ERR missing command arguments\n"}
    else
      {:reply, "ERR unknown command\n"}
    end
  end

  defp ip!(ip), do: :inet.ntoa(ip)

  defp status_loop(socket, transport, server) do
    receive do
      {:status, ^server, new} -> transport.send(socket, ["STATUS ", to_string(new), ?\n])
    end

    status_loop(socket, transport, server)
  end

  ### Cache ###

  defp check_cache(topic, opts, server, restore) do
    if opts == [] do
      server
      |> Config.agents(topic)
      |> Enum.map(& &1.agent)
      |> check_cache_integrity(topic, restore)
    else
      partition = Enum.join(opts, " ")

      server
      |> Config.agents(topic)
      |> Enum.filter(&(&1.partition == partition))
      |> Enum.map(& &1.agent)
      |> check_cache_integrity(topic, restore)
    end
  end

  defp check_cache_integrity(agents, topic, restore)
  defp check_cache_integrity([], _topic, _restore), do: {:reply, "ERR unknown_agent\n"}

  defp check_cache_integrity(agents, topic, restore) do
    {:async, "OK\n",
     fn sender ->
       total =
         agents
         |> Enum.map(&Task.async(&1, :check_cache_integrity, [restore]))
         |> Enum.map(&Task.await(&1, :infinity))
         |> Enum.sum()

       sender.("VERIFY_CACHE #{topic} #{total}\n")
     end}
  end

  ### Partition Matching ###
  @spec match_partition?(String.t(), term) :: boolean
  defp match_partition?(partition, key)
  defp match_partition?("*", _key), do: true

  defp match_partition?(partition, key) do
    [match, compare, value] = String.split(partition, " ")

    match = match_partition_match(match, key)
    value = match_partition_value(value)

    match_partition_compare(compare, match, value)
  end

  defp match_partition_match(".", key), do: key

  defp match_partition_match(pattern, key) do
    pattern
    |> String.split(".", trim: true)
    |> Enum.reduce(key, &Map.get(&2, String.to_atom(&1)))
  end

  defp match_partition_compare("=", match, value), do: match == value
  defp match_partition_compare("<", match, value), do: match < value
  defp match_partition_compare("<=", match, value), do: match <= value
  defp match_partition_compare(">", match, value), do: match > value
  defp match_partition_compare(">=", match, value), do: match >= value

  defp match_partition_value(":" <> atom), do: String.to_atom(atom)
  defp match_partition_value("\"" <> string), do: String.slice(string, 0..-2)
  defp match_partition_value("true"), do: true
  defp match_partition_value("false"), do: false
  defp match_partition_value(value), do: String.to_integer(value)
end