zen_websocket

defmodule ZenWebsocket.RequestCorrelator do
  @moduledoc """
  Manages request/response correlation for JSON-RPC WebSocket connections.

  Pure functional module - state ownership stays with Client GenServer.
  Tracks pending requests with timeouts and matches responses by ID.

  ## Telemetry Events

  The following telemetry events are emitted:

  * `[:zen_websocket, :request_correlator, :track]` - Emitted when a request is tracked.
    * Measurements: `%{count: 1}`
    * Metadata: `%{id: id, timeout_ms: timeout}`

  * `[:zen_websocket, :request_correlator, :resolve]` - Emitted when a response is matched.
    * Measurements: `%{count: 1, round_trip_ms: milliseconds}`
    * Metadata: `%{id: id}`

  * `[:zen_websocket, :request_correlator, :timeout]` - Emitted when a request times out.
    * Measurements: `%{count: 1}`
    * Metadata: `%{id: id}`

  * `[:zen_websocket, :request_correlator, :fail_all]` - Emitted for each pending request
    failed via `fail_all/2` (e.g., on connection loss).
    * Measurements: `%{count: 1}`
    * Metadata: `%{id: id, reason: reason}`
  """

  use Descripex, namespace: "/correlation"

  @typedoc "Client state map containing pending_requests field (subset of Client.state)"
  @type state :: %{
          :pending_requests => %{optional(term()) => {GenServer.from(), reference(), integer()}},
          optional(atom()) => term()
        }

  api(:extract_id, "Extract the request ID from a JSON message.",
    params: [message: [kind: :value, description: "Raw JSON binary message"]],
    returns: %{type: "{:ok, term()} | :no_id", description: "Extracted ID or :no_id if absent/invalid"}
  )

  @doc """
  Extracts the request ID from a JSON message.

  Returns `{:ok, id}` if the message contains a non-nil ID field,
  or `:no_id` if no ID is present, the message is not valid JSON,
  or the ID is nil.
  """
  @spec extract_id(binary()) :: {:ok, term()} | :no_id
  def extract_id(message) when is_binary(message) do
    case Jason.decode(message) do
      {:ok, %{"id" => id}} when not is_nil(id) -> {:ok, id}
      _ -> :no_id
    end
  end

  def extract_id(_), do: :no_id

  api(:track, "Track a pending request with a timeout timer.",
    params: [
      state: [kind: :value, description: "Client state containing pending_requests"],
      id: [kind: :value, description: "Request ID to track"],
      from: [kind: :value, description: "GenServer caller reference"],
      timeout_ms: [kind: :value, description: "Timeout in milliseconds"]
    ],
    returns: %{
      type: "{:ok, state()} | {:error, :duplicate_id, state()}",
      description: "Updated state on success, or unchanged state on duplicate ID collision"
    }
  )

  @doc """
  Tracks a pending request with a timeout timer.

  Creates a timer that will send `{:timeout, timeout_ref, {:correlation_timeout, id}}`
  to `self()` after the specified timeout. Must be called from within a
  GenServer context.

  Returns `{:ok, new_state}` on success. If `id` already has a pending entry,
  returns `{:error, :duplicate_id, state}` with `state` unchanged — the
  existing caller's `from` and timeout timer are preserved. No timer is
  started and no `:track` telemetry is emitted on the duplicate path.
  """
  @spec track(state(), term(), GenServer.from(), pos_integer()) ::
          {:ok, state()} | {:error, :duplicate_id, state()}
  def track(state, id, from, timeout_ms) do
    if Map.has_key?(state.pending_requests, id) do
      {:error, :duplicate_id, state}
    else
      timeout_ref = :erlang.start_timer(timeout_ms, self(), {:correlation_timeout, id})
      start_time = System.monotonic_time(:millisecond)
      pending = Map.put(state.pending_requests, id, {from, timeout_ref, start_time})

      :telemetry.execute(
        [:zen_websocket, :request_correlator, :track],
        %{count: 1},
        %{id: id, timeout_ms: timeout_ms}
      )

      {:ok, %{state | pending_requests: pending}}
    end
  end

  api(:resolve, "Resolve a pending request by ID, returning the caller info.",
    params: [
      state: [kind: :value, description: "Client state containing pending_requests"],
      id: [kind: :value, description: "Request ID to resolve"]
    ],
    returns: %{
      type: "{{GenServer.from(), reference(), integer()} | nil, state()}",
      description: "Entry tuple (or nil) and updated state"
    }
  )

  @doc """
  Resolves a pending request by ID, returning the caller info.

  Cancels the timeout timer and removes the request from pending.
  Returns `{entry, new_state}` where entry is `{from, timeout_ref, start_time}` or `nil`.
  """
  @spec resolve(state(), term()) :: {{GenServer.from(), reference(), integer()} | nil, state()}
  def resolve(state, id) do
    case Map.pop(state.pending_requests, id) do
      {nil, _} ->
        {nil, state}

      {{_from, timeout_ref, start_time} = entry, new_pending} ->
        Process.cancel_timer(timeout_ref, async: false, info: false)
        round_trip_ms = System.monotonic_time(:millisecond) - start_time

        :telemetry.execute(
          [:zen_websocket, :request_correlator, :resolve],
          %{count: 1, round_trip_ms: round_trip_ms},
          %{id: id}
        )

        {entry, %{state | pending_requests: new_pending}}
    end
  end

  api(:timeout, "Handle a timeout for a pending request.",
    params: [
      state: [kind: :value, description: "Client state containing pending_requests"],
      id: [kind: :value, description: "Request ID that timed out"]
    ],
    returns: %{
      type: "{{GenServer.from(), reference(), integer()} | nil, state()}",
      description: "Entry tuple (or nil) and updated state"
    }
  )

  @doc """
  Handles a timeout for a pending request.

  Removes the request from pending and returns the caller info.
  Returns `{entry, new_state}` where entry is `{from, timeout_ref, start_time}` or `nil`.
  """
  @spec timeout(state(), term()) :: {{GenServer.from(), reference(), integer()} | nil, state()}
  def timeout(state, id) do
    case Map.pop(state.pending_requests, id) do
      {nil, _} ->
        {nil, state}

      {entry, new_pending} ->
        :telemetry.execute(
          [:zen_websocket, :request_correlator, :timeout],
          %{count: 1},
          %{id: id}
        )

        {entry, %{state | pending_requests: new_pending}}
    end
  end

  api(:fail_all, "Fail every pending request with the given reason.",
    params: [
      state: [kind: :value, description: "Client state containing pending_requests"],
      reason: [kind: :value, description: "Error reason delivered to each caller (e.g., :disconnected)"]
    ],
    returns: %{type: "state()", description: "State with pending_requests cleared"}
  )

  @doc """
  Fails every pending request with `{:error, reason}` and cancels their timeout timers.

  Used when the connection drops and in-flight correlated responses can no longer
  arrive. Callers blocked on `GenServer.call` receive the reply immediately
  instead of waiting for their per-call timeout.
  """
  @spec fail_all(state(), term()) :: state()
  def fail_all(state, reason) do
    Enum.each(state.pending_requests, fn {id, {from, timeout_ref, _start_time}} ->
      Process.cancel_timer(timeout_ref, async: false, info: false)
      GenServer.reply(from, {:error, reason})

      :telemetry.execute(
        [:zen_websocket, :request_correlator, :fail_all],
        %{count: 1},
        %{id: id, reason: reason}
      )
    end)

    %{state | pending_requests: %{}}
  end

  api(:pending_count, "Return the count of pending requests.",
    params: [state: [kind: :value, description: "Client state containing pending_requests"]],
    returns: %{type: "non_neg_integer()", description: "Number of requests awaiting response"}
  )

  @doc """
  Returns the count of pending requests.
  """
  @spec pending_count(state()) :: non_neg_integer()
  def pending_count(state) do
    map_size(state.pending_requests)
  end
end