defmodule Maelstrom do
  @moduledoc """
  Allows you to create servers which implement the [Maelstrom protocol](https://github.com/jepsen-io/maelstrom/blob/main/doc/protocol.md).

  [Maelstrom](https://github.com/jepsen-io/maelstrom) is a workbench for learning distributed systems by writing your own.

  # Usage

  To implement a server:

  1. Create a module and `use Maelstrom`
  2. Implement one or more `handle_message` heads.
  3. Call `MyModule.run_forever()`.

  I recommend you do this in a .exs script. Example:

  ```
  defmodule Echo do
    use Maelstrom

    def handle_message(_src, _dest, %{"echo" => echo}, state, _) do
      {:reply, %{"type" => "echo_ok", "echo" => echo}, state}
    end
  end

  Echo.run_forever()
  ```

  You could then run this with `mix run` e.g. `mix run echo.exs`.

  > #### Tip {: .tip}
  > Maelstrom expects a single binary with no arguments to call for testing.
  > In order to accomplish this, wrap your mix run command in a shell script (see [demos](https://github.com/prehnRA/maelstrom_ex/tree/main/demo) for examples).

  # Examples

  For more examples, see [demos](https://github.com/prehnRA/maelstrom_ex/tree/main/demo).
  """

  @type node_id() :: String.t()
  @type msg_id() :: non_neg_integer()
  @type error_code() :: non_neg_integer()
  @type body() :: Map
  @type state() :: Map
  @type node_state() :: Map
  @type handler_result() :: {:reply, body(), state()} | {:noreply, state()} | {:error, non_neg_integer(), String.t(), state()}
  @type rpc_callback() :: (body(), body(), state() -> handler_result())

  @doc "Docs on a callback?"
  @callback handle_message(node_id(), node_id(), body(), state(), node_state()) :: handler_result()

  @doc """
  Run the server forever and prevent the script from exiting. Handles IO to and
  from your server node. You should invoke this directly on your server module
  instead of calling it on `Maelstrom` e.g. `MyModule.run_forever()`. This function
  is defined automatically when you `use Maelstrom`.
  """
  @spec run_forever(state()) :: :ok
  def run_forever(_initial_state \\ %{}) do
    # Stub for documentation
    raise "Call this function directly on your server module e.g. MyModule.run_forever()."
  end

  @doc false
  def take_input_forever(pid) do
    stream = IO.stream()
    Task.async(fn ->
      Enum.each(stream, fn line ->
        message = Jason.decode!(line)
        log("Received: #{inspect(message)}")

        GenServer.cast(pid, message)
      end)
    end)
  end

  @doc "Logs an error to :stderr in the way that Maelstrom expects."
  @spec log(String.t(), atom()) :: :ok
  def log(log, adapter \\ GenServer) do
    adapter.cast(self(), {:log, log})
  end

  @doc """
  Send an RPC-style message with the body `body` to `dest`. `callback` is a function
  that will be invoked if and when `dest` replies to the message. The arguments to callback are:

  1. The body of the original message.
  2. The body of the reply.
  3. The state of your server.

  The callback should reply with a valid handler response:

  1. `{:reply, reply_body, new_server_state}`
  2. `{:noreply, new_server_state}`
  3. `{:error, error_code, error_message, new_server_state}`
  """
  @spec send_rpc(node_id(), body(), rpc_callback(), atom()) :: :ok
  def send_rpc(dest, body, callback, adapter \\ GenServer) do
    adapter.cast(self(), {:send_rpc, dest, body, callback})
  end

  @doc """
  Send a message with the body `body` to `dest`.
  """
  @spec send_message(node_id(), body(), atom()) :: :ok
  def send_message(dest, body, adapter \\ GenServer) do
    adapter.cast(self(), {:send_message, dest, body})
  end

  @doc """
  Send a reply to `replyee` replying to the message with id `in_reply_to`. Include the body `body`.
  """
  @spec send_reply(node_id(), msg_id(), body(), atom()) :: :ok
  def send_reply(replyee, in_reply_to, body, adapter \\ GenServer) do
    send_message(replyee, Map.put(body, "in_reply_to", in_reply_to), adapter)
  end

  @doc """
  Send an error message to `replyee` in reply to the message indentified by `in_reply_to`. The error
  message will include the error code `error_code` (see
  [the Maelstrom docs](https://github.com/jepsen-io/maelstrom/blob/main/doc/protocol.md#errors) for
  a list of codes) and a string message `error_text`.
  """
  @spec send_error(node_id(), msg_id(), error_code(), String.t(), atom()) :: :ok
  def send_error(replyee, in_reply_to, error_code, error_text, adapter \\ GenServer) do
    body =
      %{
        "type" => "error",
        "code" => error_code,
        "text" => error_text
      }

    send_reply(replyee, in_reply_to, body, adapter)
  end

  defmacro __using__(_opts \\ []) do
    quote do
      use GenServer

      require Logger
      import Maelstrom

      unquote(setup_cast_handlers())
    end
  end

  @doc false
  def default_node_state do
    %{msg_id: 1, rpcs: %{}}
  end

  @doc false
  def setup_cast_handlers() do
    quote do
      def run_forever(initial_state \\ %{}) do
        {:ok, pid} = GenServer.start_link(__MODULE__, initial_state)
        take_input_forever(pid) |> Task.await(:infinity)
      end

      @impl true
      def init(inner_state \\ %{}) do
        {:ok, {Maelstrom.default_node_state(), inner_state}}
      end

      @impl true
      def handle_cast(%{
        "src" => src,
        "body" => %{
          "type" => "init",
          "msg_id" => msg_id,
          "node_id" => node_id,
          "node_ids" => node_ids
        }}, {node, inner_state}) do

        new_node_state = Map.merge(node, %{node_id: node_id, node_ids: node_ids})

        send_reply(src, msg_id, %{"type" => "init_ok"})
        {:noreply, {new_node_state, inner_state}}
      end

      def handle_cast(%{
        "src" => src,
        "body" => %{
          "in_reply_to" => msg_id
        } = reply_body} = msg, {%{rpcs: rpcs} = node_state, inner_state} = state) do

        {original_body, callback} = Map.get(rpcs, msg_id, {nil, &warn_missing_rpc_callback/3})
        process_callback(callback.(original_body, reply_body, inner_state), msg, state)
      end

      @impl true
      def handle_cast(%{
        "src" => src,
        "dest" => dest,
        "body" => %{"msg_id" => msg_id} = body
      } = msg, {node_state, inner_state} = state) do
        process_callback(handle_message(src, dest, body, inner_state, node_state), msg, state)
      end

      def handle_cast({:send_rpc, dest, body, callback}, {%{rpcs: rpcs, msg_id: msg_id} = node_state, inner_state}) do

        {:noreply, {new_node_state, inner_state}} = handle_cast({:send_message, dest, body}, {node_state, inner_state})

        rpcs = Map.put(rpcs, msg_id, {body, callback})

        node_state = Map.put(node_state, :rpcs, rpcs)

        {:noreply, {node_state, inner_state}}
      end

      def handle_cast({:send_message, dest, body}, {%{node_id: src, msg_id: msg_id} = node_state, inner_state}) do
         message = %{
          "src" => src,
          "dest" => dest,
          "body" => Map.put_new(body, "msg_id", msg_id)
        }
        json = Jason.encode!(message)

        IO.puts(json)

        new_node_state = Map.put(node_state, :msg_id, msg_id + 1)

        {:noreply, {new_node_state, inner_state}}
      end

      def handle_cast({:log, log}, state) do
        IO.warn(log)

        {:noreply, state}
      end

      def handle_cast(other, _state) do
        raise "Unexpected message format: #{inspect(other)}"
      end

      defp process_callback(result, %{
        "src" => src,
        "body" => %{"msg_id" => msg_id} = body
      }, {node_state, _}) do
        case result do
          {:reply, body, new_inner_state} ->
            send_reply(src, msg_id, body)
            {:noreply, {node_state, new_inner_state}}

          {:noreply, new_inner_state} ->
            {:noreply, {node_state, new_inner_state}}

          {:error, error_code, error_text, new_inner_state} ->
            send_error(src, msg_id, error_code, error_text)
            {:noreply, {node_state, new_inner_state}}
        end
      end

      defp warn_missing_rpc_callback(_message_body, reply_body, inner_state) do
        IO.warn("Reply for unknown RPC: reply_body=#{inspect(reply_body)}")

        {:noreply, inner_state}
      end
    end
  end
end