mammoth

defmodule Mammoth do
  @moduledoc ~S"""
  Mammoth: A STOMP client.

  Use `Mammoth` as the primary API and `Mammoth.Message` for working with received messages.

  My recommendation is to start mammoth from within your own implementation of the callback handler,
    and to put the callback handler itself into a supervisor

  ## Example
      {:ok, callback_pid} = Mammoth.DefaultCallbackHandler.start_link
      {:ok, pid} = Mammoth.start_link(callback_pid)
      :ok = Mammoth.DefaultCallbackHandler.set_mammoth_pid(callback_pid, pid)
      Mammoth.connect(pid, {127,0,0,1}, 61613, "/", "admin", "admin")
      Mammoth.subscribe(pid, "foo.bar", :client)
      Mammoth.disconnect(pid)

  ## Starting in a supervision tree

      children = [
        worker(Mammoth, [%{}, [name: Mammoth]])
      ]
  """

  use GenServer

  require Logger
  alias Mammoth.{Message, Receiver, Socket, Subscriber}

  def init(args) do
    {
      :ok,
      args
      |> Map.put_new(:socket, nil)
      |> Map.put_new(:receiver, nil)
      |> Map.put_new(:subscriber, nil)
    }
  end

  def start_link(callback_handler, state \\ %{}, opts \\ []) do
    GenServer.start_link(
      __MODULE__,
      state |> Map.put_new(:callback_handler, callback_handler),
      opts
    )
  end

  @doc """
  Connect to server.

  `host` must be `inet:socket_address() | inet:hostname()`, for example `{127,0,0,1}` or `'example.com'`.
  `virtual_host` (`host` header) is mandatory in STOMP 1.2, behaviour is server defined. Default on RabbitMQ would be `"/"`.

  https://stomp.github.io/stomp-specification-1.2.html#CONNECT_or_STOMP_Frame
  """
  def connect(pid, host, port, virtual_host, login, password, additional_headers \\ []) do
    GenServer.call(pid, {:connect, host, port, virtual_host, login, password, additional_headers})
  end

  @doc """
  Subscribe to a queue.

  Note that if you set ack_mode to :client_individual or :client,
  you must send ACK frames when you receive MESSAGE frames (not required with :auto)
  https://stomp.github.io/stomp-specification-1.2.html#SUBSCRIBE_ack_Header
  """
  def subscribe(pid, destination, ack_mode \\ :auto, additional_headers \\ []) do
    GenServer.call(pid, {:subscribe, destination, get_ack_mode(ack_mode), additional_headers})
  end

  defp get_ack_mode(ack_mode) do
    case ack_mode do
      :auto -> "auto"
      :client_individual -> "client-individual"
      :client -> "client"
    end
  end

  @doc """
  Unsubscribe from a queue.
  """
  def unsubscribe(pid, destination) do
    GenServer.call(pid, {:unsubscribe, destination})
  end

  @doc """
  Receive messages from the TCP socket.

  Is called automatically when necessary. Should not be called manually.
  """
  def receive(pid, message) do
    GenServer.call(pid, {:receive, message})
  end

  @doc """
  Disconnect from server.
  """
  def disconnect(pid) do
    GenServer.call(pid, :disconnect)
  end

  @doc """
  Send a frame to the remote server
  """
  def send_frame(pid, message = %Message{}) do
    GenServer.call(pid, {:send, message})
  end

  @doc """
  Send a SEND message to the remote server

  content-length header will be automatically added.
  """
  def send_send(pid, destination, body, additional_headers \\ []) do
    GenServer.call(
      pid,
      {:send,
       %Message{
         command: :send,
         headers: [{"destination", destination}] ++ additional_headers,
         body: body
       }}
    )
  end

  @doc """
  Send an ACK message to the remote server for the specified frame
  """
  def send_ack_frame(pid, message = %Message{}, additional_headers \\ []) do
    {:ok, message_id} = Message.get_header(message, "ack")
    send_ack_id(pid, message_id, additional_headers)
  end

  @doc """
  Send an ACK message to the remote server for the specified frame ID (referenced in the MESSAGE `ack` header)
  """
  def send_ack_id(pid, id, additional_headers) do
    GenServer.call(
      pid,
      {:send, %Message{command: :ack, headers: [{"id", id}] ++ additional_headers}}
    )
  end

  @doc """
  Send an NACK message to the remote server for the specified frame
  """
  def send_nack_frame(pid, message = %Message{}, additional_headers \\ []) do
    {:ok, message_id} = Message.get_header(message, "message-id")
    send_nack_id(pid, message_id, additional_headers)
  end

  @doc """
  Send an NACK message to the remote server for the specified frame ID (referenced in the MESSAGE `ack` header)
  """
  def send_nack_id(pid, id, additional_headers) do
    GenServer.call(
      pid,
      {:send, %Message{command: :nack, headers: [{"id", id}] ++ additional_headers}}
    )
  end

  def handle_call(
        {:send, message},
        _from,
        state = %{socket: socket}
      ) do
    Socket.send(socket, message)
    {:reply, :ok, state}
  end

  def handle_call(
        {:receive, message},
        _from,
        state = %{callback_handler: callback_handler}
      ) do
    Kernel.send(callback_handler, {:mammoth, :receive_frame, message})

    {:reply, :ok, state}
  end

  def handle_call(
        {:connect, host, port, virtual_host, login, password, additional_headers},
        _from,
        state
      ) do
    # todo: handle {:error, :econnrefused} response
    {:ok, socket} = Socket.connect(host, port)
    Socket.send(socket, connect_message(virtual_host, login, password, additional_headers))

    {:ok, subscriber} = Subscriber.start_link()
    {:ok, receiver} = Receiver.start_link(%{socket: socket, consumer: self()})
    Receiver.listen(receiver)

    {:reply, socket, %{state | socket: socket, subscriber: subscriber, receiver: receiver}}
  end

  @doc """
  Requests disconnection from the remote server
  """
  def handle_call(
        :disconnect,
        _from,
        state = %{
          socket: socket
        }
      ) do
    receipt_id = Enum.random(1000..1_000_000)
    Socket.send(socket, disconnect_message(receipt_id))

    {:noreply, Map.put(state, :disconnect_id, receipt_id)}
  end

  def handle_call(
        {:subscribe, destination, ack_mode, additional_headers},
        _from,
        state = %{socket: socket, subscriber: subscriber}
      ) do
    {:ok, entry} = Subscriber.subscribe(subscriber, destination, ack_mode)
    message = subscribe_message(destination, entry.id, ack_mode, additional_headers)
    Socket.send(socket, message)
    {:reply, {:ok, entry}, state}
  end

  def handle_call(
        {:unsubscribe, destination},
        _from,
        state = %{socket: socket, subscriber: subscriber}
      ) do
    {:ok, entry} = Subscriber.unsubscribe(subscriber, destination)
    message = unsubscribe_message(entry.id)
    Socket.send(socket, message)
    {:reply, :ok, state}
  end

  def handle_cast(
        :disconnected,
        state = %{
          subscriber: subscriber,
          receiver: receiver,
          callback_handler: callback_handler,
          disconnect_id: _disconnect_id
        }
      ) do
    Subscriber.stop(subscriber)
    Receiver.stop(receiver)

    Kernel.send(callback_handler, {:mammoth, :disconnected, true})

    {:noreply, state}
  end

  def handle_cast(
        :disconnected,
        state = %{
          subscriber: subscriber,
          receiver: receiver,
          callback_handler: callback_handler
        }
      ) do
    Subscriber.stop(subscriber)
    Receiver.stop(receiver)

    Kernel.send(callback_handler, {:mammoth, :disconnected, false})

    {:noreply, state}
  end

  defp connect_message(virtual_host, login, password, additional_headers) do
    %Message{
      command: :connect,
      headers:
        [
          {"accept-version", "1.2"},
          {"host", virtual_host},
          {"login", login},
          {"passcode", password}
        ] ++ additional_headers
    }
  end

  defp disconnect_message(receipt_id) do
    %Message{
      command: :disconnect,
      headers: [
        {"receipt-id", receipt_id}
      ]
    }
  end

  defp subscribe_message(destination, id, ack_mode, additional_headers) do
    %Message{
      command: :subscribe,
      headers:
        [
          {"destination", destination},
          {"ack", ack_mode},
          {"id", id}
        ] ++ additional_headers
    }
  end

  defp unsubscribe_message(id) do
    %Message{
      command: :unsubscribe,
      headers: [
        {"id", id}
      ]
    }
  end
end