elixir_stomp_client

defmodule StompClient do
  @moduledoc """
  Custom STOMP protocol client implementation for Elixir.

  This module provides a GenServer-based STOMP client with precise control over message formatting, 
  especially content-length headers. It features automatic reconnection, heartbeat support, and 
  callback-based message handling.

  ## Features

  - **Custom Message Formatting**: Control over content-length headers (can omit content-length for SEND frames)
  - **Automatic Reconnection**: Automatic reconnection with configurable retry intervals
  - **Heartbeat Support**: Built-in heartbeat mechanism to keep connections alive
  - **Callback-based Handling**: Process-based callback system for handling connection events and messages
  - **STOMP 1.2 Compatible**: Full support for STOMP 1.2 protocol
  - **Subscription Management**: Easy subscription and unsubscription management

  ## Usage

  The basic workflow involves:

  1. Start a callback handler process
  2. Connect to the STOMP server
  3. Subscribe to destinations
  4. Send/receive messages
  5. Disconnect when done

  ### Example

      # Define a callback handler
      defmodule MyHandler do
        use GenServer
        
        def start_link(_), do: GenServer.start_link(__MODULE__, nil, name: __MODULE__)
        def init(_), do: {:ok, %{}}
        
        def handle_info({:stomp_client, :on_connect, _frame}, state) do
          IO.puts("Connected!")
          {:noreply, state}
        end
        
        def handle_info({:stomp_client, :on_message, message}, state) do
          IO.puts("Received message")
          {:noreply, state}
        end
        
        def handle_info(_, state), do: {:noreply, state}
      end

      # Start handler and connect  
      {:ok, _} = MyHandler.start_link(nil)
      {:ok, client} = StompClient.connect([
        host: "localhost",
        port: 61613,
        username: "guest",
        password: "guest"
      ], callback_handler: MyHandler)

      # Use the client
      StompClient.subscribe(client, "/queue/test")
      StompClient.send(client, "/queue/test", "Hello World!")
      StompClient.disconnect(client)

  ## Callback Messages

  Your callback handler will receive these message types:

  - `{:stomp_client, :on_connect, frame}` - When successfully connected
  - `{:stomp_client, :on_message, message}` - When a message is received  
  - `{:stomp_client, :on_disconnect, nil}` - When disconnected
  - `{:stomp_client, :on_connect_error, error}` - When connection fails
  - `{:stomp_client, :on_send, frame}` - When a frame is sent (optional)
  """

  use GenServer
  require Logger

  # STOMP protocol constants
  @stomp_version "1.2"
  @null_byte "\0"
  @newline "\n"

  defstruct [
    :socket,
    :host,
    :port,
    :username,
    :password,
    :connected,
    :session_id,
    :subscriptions,
    :callback_handler,
    :heartbeat_timer,
    :receive_buffer
  ]

  # Client API

  @doc """
  Starts a STOMP client process.

  ## Parameters

  - `opts` - Connection options as keyword list or map

  ## Options

  - `:host` - STOMP server hostname (required)
  - `:port` - STOMP server port (default: 61613)
  - `:username` or `:login` - Authentication username  
  - `:password` or `:passcode` - Authentication password
  - `:callback_handler` - Process to receive callback messages (required)

  ## Examples

      {:ok, client} = StompClient.start_link([
        host: "localhost",
        port: 61613,
        username: "guest",
        password: "guest",
        callback_handler: MyHandler
      ])

  """
  @spec start_link(keyword() | map()) :: GenServer.on_start()
  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts)
  end

  @doc """
  Connects to a STOMP server with callback handler.

  This is a convenience function that combines `start_link/1` with callback handler setup.

  ## Parameters

  - `opts` - Connection options (keyword list or map)
  - `callback_handler` - Process to receive callback messages

  ## Examples

      {:ok, client} = StompClient.connect([
        host: "localhost",
        port: 61613,
        username: "guest", 
        password: "guest"
      ], callback_handler: MyHandler)

  """
  @spec connect(keyword() | map(), callback_handler: pid() | atom()) :: GenServer.on_start()
  def connect(opts, callback_handler: handler) when is_list(opts) do
    opts_with_handler = Keyword.put(opts, :callback_handler, handler)
    GenServer.start_link(__MODULE__, opts_with_handler)
  end

  def connect(opts, callback_handler: handler) when is_map(opts) do
    opts_with_handler = Map.put(opts, :callback_handler, handler)
    GenServer.start_link(__MODULE__, opts_with_handler)
  end

  @doc """
  Sends a message to a STOMP destination.

  ## Parameters

  - `pid` - The STOMP client process
  - `destination` - STOMP destination (e.g., "/queue/test" or "/topic/news")
  - `body` - Message body as string
  - `headers` - Optional list of additional headers as tuples

  ## Returns

  - `:ok` on success
  - `{:error, :not_connected}` if client is not connected
  - `{:error, reason}` on other errors

  ## Examples

      # Simple message
      StompClient.send(client, "/queue/test", "Hello World!")
      
      # Message with custom headers  
      StompClient.send(client, "/queue/test", "Hello", [
        {"priority", "5"},
        {"custom-header", "value"}
      ])

  """
  @spec send(pid(), String.t(), String.t(), [{String.t(), String.t()}]) :: 
    :ok | {:error, :not_connected | term()}
  def send(pid, destination, body, headers \\ []) do
    GenServer.call(pid, {:send, destination, body, headers})
  end

  @doc """
  Subscribes to a STOMP destination.

  ## Parameters

  - `pid` - The STOMP client process
  - `destination` - STOMP destination to subscribe to
  - `opts` - Optional subscription options

  ## Options

  - `:id` - Custom subscription ID (auto-generated if not provided)
  - `:selector` - Message selector for filtering

  ## Returns

  - `:ok` on success
  - `{:error, :not_connected}` if client is not connected
  - `{:error, reason}` on other errors

  ## Examples

      # Simple subscription
      StompClient.subscribe(client, "/queue/test")
      
      # Subscription with selector
      StompClient.subscribe(client, "/queue/test", selector: "priority > 5")
      
      # Subscription with custom ID
      StompClient.subscribe(client, "/queue/test", id: "my-subscription")

  """
  @spec subscribe(pid(), String.t(), keyword()) :: 
    :ok | {:error, :not_connected | term()}
  def subscribe(pid, destination, opts \\ []) do
    GenServer.call(pid, {:subscribe, destination, opts})
  end

  @doc """
  Unsubscribes from a STOMP destination.

  ## Parameters

  - `pid` - The STOMP client process
  - `destination` - STOMP destination to unsubscribe from

  ## Returns

  - `:ok` on success
  - `{:error, :not_connected}` if client is not connected  
  - `{:error, :not_subscribed}` if not subscribed to destination
  - `{:error, reason}` on other errors

  ## Examples

      StompClient.unsubscribe(client, "/queue/test")

  """
  @spec unsubscribe(pid(), String.t()) :: 
    :ok | {:error, :not_connected | :not_subscribed | term()}
  def unsubscribe(pid, destination) do
    GenServer.call(pid, {:unsubscribe, destination})
  end

  @doc """
  Disconnects from the STOMP server and stops the client.

  ## Parameters

  - `pid` - The STOMP client process

  ## Returns

  - `:ok` always

  ## Examples

      StompClient.disconnect(client)

  """
  @spec disconnect(pid()) :: :ok
  def disconnect(pid) do
    GenServer.call(pid, :disconnect)
  end

  # GenServer Callbacks

  @impl true
  def init(opts) when is_list(opts) do
    opts_map = Enum.into(opts, %{})
    init(opts_map)
  end

  def init(opts) when is_map(opts) do
    state = %__MODULE__{
      host: (opts[:host] || opts["host"]) |> to_string(),
      port: opts[:port] || opts["port"] || 61613,
      username: opts[:login] || opts[:username] || opts["login"] || opts["username"],
      password: opts[:passcode] || opts[:password] || opts["passcode"] || opts["password"],
      connected: false,
      subscriptions: %{},
      callback_handler: opts[:callback_handler] || opts["callback_handler"],
      receive_buffer: ""
    }

    send(self(), :connect)
    {:ok, state}
  end

  @impl true
  def handle_info(:connect, state) do
    case establish_connection(state) do
      {:ok, new_state} ->
        send_connect_frame(new_state)
        {:noreply, new_state}

      {:error, reason} ->
        Logger.error("Failed to connect: #{inspect(reason)}")
        Process.send_after(self(), :connect, 5000)
        {:noreply, state}
    end
  end

  @impl true
  def handle_info({:tcp, socket, data}, %{socket: socket} = state) do
    Logger.debug("Received TCP data: #{inspect(data)}")
    new_buffer = state.receive_buffer <> data
    {frames, remaining_buffer} = parse_frames(new_buffer)
    
    updated_state = %{state | receive_buffer: remaining_buffer}
    
    final_state = Enum.reduce(frames, updated_state, fn frame, acc_state ->
      Logger.debug("Parsed STOMP frame: #{inspect(frame)}")
      handle_stomp_frame(frame, acc_state)
    end)
    
    {:noreply, final_state}
  end

  @impl true
  def handle_info({:tcp_closed, socket}, %{socket: socket} = state) do
    Logger.warning("STOMP connection closed")
    notify_handler({:stomp_client, :on_disconnect, nil}, state)
    cancel_heartbeat_timer(state)
    Process.send_after(self(), :connect, 5000)
    {:noreply, %{state | socket: nil, connected: false, heartbeat_timer: nil}}
  end

  @impl true
  def handle_info({:tcp_error, socket, reason}, %{socket: socket} = state) do
    Logger.error("STOMP TCP error: #{inspect(reason)}")
    notify_handler({:stomp_client, :on_connect_error, %{"message" => inspect(reason)}}, state)
    cancel_heartbeat_timer(state)
    Process.send_after(self(), :connect, 5000)
    {:noreply, %{state | socket: nil, connected: false, heartbeat_timer: nil}}
  end

  @impl true
  def handle_info(:send_heartbeat, state) do
    if state.connected and state.socket do
      Logger.debug("Sending heartbeat")
      send_heartbeat_frame(state)
      # Schedule next heartbeat
      heartbeat_timer = Process.send_after(self(), :send_heartbeat, 25_000)
      {:noreply, %{state | heartbeat_timer: heartbeat_timer}}
    else
      {:noreply, state}
    end
  end

  @impl true
  def handle_call({:send, destination, body, headers}, _from, state) do
    if state.connected do
      result = send_message_frame(state, destination, body, headers)
      {:reply, result, state}
    else
      {:reply, {:error, :not_connected}, state}
    end
  end

  @impl true
  def handle_call({:subscribe, destination, opts}, _from, state) do
    if state.connected do
      id = Keyword.get(opts, :id, generate_subscription_id())
      selector = Keyword.get(opts, :selector)
      
      result = send_subscribe_frame(state, destination, id, selector)
      
      case result do
        :ok ->
          new_subscriptions = Map.put(state.subscriptions, id, destination)
          {:reply, :ok, %{state | subscriptions: new_subscriptions}}
        error ->
          {:reply, error, state}
      end
    else
      {:reply, {:error, :not_connected}, state}
    end
  end

  @impl true
  def handle_call({:unsubscribe, destination}, _from, state) do
    if state.connected do
      # Find subscription ID for destination
      case find_subscription_id(state.subscriptions, destination) do
        {:ok, id} ->
          result = send_unsubscribe_frame(state, id)
          new_subscriptions = Map.delete(state.subscriptions, id)
          {:reply, result, %{state | subscriptions: new_subscriptions}}
        :error ->
          {:reply, {:error, :not_subscribed}, state}
      end
    else
      {:reply, {:error, :not_connected}, state}
    end
  end

  @impl true
  def handle_call(:disconnect, _from, state) do
    if state.socket do
      send_disconnect_frame(state)
      :gen_tcp.close(state.socket)
    end
    cancel_heartbeat_timer(state)
    {:reply, :ok, %{state | socket: nil, connected: false, heartbeat_timer: nil}}
  end

  # Private Functions

  defp establish_connection(state) do
    case :gen_tcp.connect(
      String.to_charlist(state.host),
      state.port,
      [:binary, packet: :raw, active: true, keepalive: true]
    ) do
      {:ok, socket} ->
        {:ok, %{state | socket: socket}}
      {:error, reason} ->
        {:error, reason}
    end
  end

  defp send_connect_frame(state) do
    headers = [
      {"accept-version", @stomp_version},
      {"host", state.host},
      {"login", state.username},
      {"passcode", state.password},
      {"heart-beat", "30000,30000"}
    ]

    frame = build_frame("CONNECT", headers, "")
    Logger.info("Sending CONNECT frame to #{state.host}:#{state.port}")
    send_frame(state, frame)
  end

  defp send_message_frame(state, destination, body, custom_headers) do
    # This is where we control the content-length behavior
    # We explicitly DO NOT include content-length to match Go's NoContentLength
    headers = [
      {"destination", destination},
      {"content-type", "text/plain"}
    ] ++ custom_headers

    frame = build_frame_without_content_length("SEND", headers, body)
    send_frame(state, frame)
  end

  defp send_subscribe_frame(state, destination, id, selector) do
    headers = [
      {"destination", destination},
      {"id", id},
      {"ack", "auto"}
    ]

    headers = if selector do
      [{"selector", selector} | headers]
    else
      headers
    end

    frame = build_frame("SUBSCRIBE", headers, "")
    send_frame(state, frame)
  end

  defp send_unsubscribe_frame(state, id) do
    headers = [{"id", id}]
    frame = build_frame("UNSUBSCRIBE", headers, "")
    send_frame(state, frame)
  end

  defp send_disconnect_frame(state) do
    frame = build_frame("DISCONNECT", [], "")
    send_frame(state, frame)
  end

  defp send_heartbeat_frame(state) do
    # STOMP heartbeat is just a newline character
    case :gen_tcp.send(state.socket, @newline) do
      :ok -> :ok
      {:error, reason} -> 
        Logger.error("Failed to send heartbeat: #{inspect(reason)}")
        {:error, reason}
    end
  end

  defp cancel_heartbeat_timer(state) do
    if state.heartbeat_timer do
      Process.cancel_timer(state.heartbeat_timer)
    end
  end

  # Standard STOMP frame with content-length
  defp build_frame(command, headers, body) do
    body_length = byte_size(body)
    all_headers = [{"content-length", to_string(body_length)} | headers]
    
    header_lines = Enum.map(all_headers, fn {key, value} -> "#{key}:#{value}" end)
    
    command <> @newline <>
    Enum.join(header_lines, @newline) <> @newline <>
    @newline <> body <> @null_byte
  end

  # STOMP frame WITHOUT content-length (like Go's NoContentLength)
  defp build_frame_without_content_length(command, headers, body) do
    header_lines = Enum.map(headers, fn {key, value} -> "#{key}:#{value}" end)
    
    command <> @newline <>
    Enum.join(header_lines, @newline) <> @newline <>
    @newline <> body <> @null_byte
  end

  defp send_frame(state, frame) do
    case :gen_tcp.send(state.socket, frame) do
      :ok -> 
        notify_handler({:stomp_client, :on_send, frame}, state)
        :ok
      {:error, reason} -> 
        {:error, reason}
    end
  end

  defp parse_frames(buffer) do
    parse_frames_recursive(buffer, [])
  end

  defp parse_frames_recursive(buffer, acc) do
    case String.split(buffer, @null_byte, parts: 2) do
      [complete_frame, remaining] ->
        # Trim any leading/trailing whitespace from the frame
        trimmed_frame = String.trim(complete_frame)
        case parse_single_frame(trimmed_frame) do
          {:ok, frame} ->
            parse_frames_recursive(remaining, [frame | acc])
          :incomplete ->
            {Enum.reverse(acc), buffer}
        end
      [incomplete] ->
        {Enum.reverse(acc), incomplete}
    end
  end

  defp parse_single_frame(frame_data) do
    case String.split(frame_data, @newline <> @newline, parts: 2) do
      [header_section, body] ->
        lines = String.split(header_section, @newline)
        command = List.first(lines) || ""
        header_lines = List.delete_at(lines, 0)
        headers = parse_headers(header_lines)
        
        Logger.debug("Parsed command: '#{command}', headers: #{inspect(headers)}")
        
        {:ok, %{
          command: command,
          headers: headers,
          body: body
        }}
      [single_section] ->
        # Handle case where there might not be a body (like heartbeats)
        lines = String.split(single_section, @newline)
        command = List.first(lines) || ""
        header_lines = List.delete_at(lines, 0)
        headers = parse_headers(header_lines)
        
        {:ok, %{
          command: command,
          headers: headers,
          body: ""
        }}
      _ ->
        :incomplete
    end
  end

  defp parse_headers(header_lines) do
    Enum.reduce(header_lines, %{}, fn line, acc ->
      case String.split(line, ":", parts: 2) do
        [key, value] -> Map.put(acc, key, value)
        _ -> acc
      end
    end)
  end

  defp handle_stomp_frame(%{command: "CONNECTED"} = frame, state) do
    Logger.info("STOMP CONNECTED frame received, connection established")
    session_id = Map.get(frame.headers, "session")
    
    # Start heartbeat timer to send heartbeats every 25 seconds (less than 30s timeout)
    heartbeat_timer = Process.send_after(self(), :send_heartbeat, 25_000)
    
    new_state = %{state | connected: true, session_id: session_id, heartbeat_timer: heartbeat_timer}
    notify_handler({:stomp_client, :on_connect, frame}, state)
    new_state
  end

  defp handle_stomp_frame(%{command: "MESSAGE"} = frame, state) do
    message_data = %{
      "destination" => Map.get(frame.headers, "destination"),
      "body" => frame.body
    }
    |> Map.merge(frame.headers)
    
    notify_handler({:stomp_client, :on_message, message_data}, state)
    state
  end

  defp handle_stomp_frame(%{command: "ERROR"} = frame, state) do
    error_data = %{
      "message" => frame.body,
      "headers" => frame.headers
    }
    notify_handler({:stomp_client, :on_message_error, error_data}, state)
    state
  end

  defp handle_stomp_frame(%{command: "RECEIPT"} = frame, state) do
    Logger.debug("Received RECEIPT: #{inspect(frame)}")
    state
  end

  defp handle_stomp_frame(%{command: ""} = frame, state) do
    # This is likely a MESSAGE frame that wasn't parsed correctly
    Logger.debug("Handling frame with empty command as MESSAGE: #{inspect(frame)}")
    message_data = %{
      "destination" => Map.get(frame.headers, "destination"),
      "body" => frame.body
    }
    |> Map.merge(frame.headers)
    
    notify_handler({:stomp_client, :on_message, message_data}, state)
    state
  end

  defp handle_stomp_frame(frame, state) do
    Logger.warning("Unknown STOMP frame: #{inspect(frame)}")
    state
  end

  defp notify_handler(message, state) do
    if state.callback_handler do
      send(state.callback_handler, message)
    end
  end

  defp generate_subscription_id do
    "sub-#{:erlang.unique_integer([:positive])}"
  end

  defp find_subscription_id(subscriptions, destination) do
    case Enum.find(subscriptions, fn {_id, dest} -> dest == destination end) do
      {id, _dest} -> {:ok, id}
      nil -> :error
    end
  end
end