s2_client

defmodule S2.S2S.ReadSession do
  @moduledoc """
  Unidirectional streaming read session.

  Opens a streaming GET to `/v1/streams/{stream}/records` with `Content-Type: s2s/proto`.
  Call `next_batch/1` repeatedly to receive `ReadBatch` messages. Heartbeat frames
  (empty ReadBatch) are skipped automatically.

  ## Process affinity

  Sessions are NOT safe to share across processes. The underlying Mint connection
  delivers TCP messages to the owning process's mailbox. Creating a session in one
  process and calling `next_batch/1` from another will not work — the receiving process
  won't see the TCP data.
  """

  require Logger

  alias S2.S2S.Shared

  @typedoc "An open read session."
  @type t :: %__MODULE__{}

  defstruct [:conn, :request_ref, :owner_pid, recv_timeout: 5_000, closed: false, data: <<>>]

  @doc """
  Open a new streaming read session.

  ## Options

  Accepts the same query opts as unary read: `:seq_num`, `:count`, `:wait`, etc.

    * `:token` — Bearer token for authentication.
    * `:recv_timeout` — Timeout in milliseconds for receiving responses (default: 5000).

  Returns `{:ok, session}` on success or `{:error, reason}` on failure.
  On `Mint.HTTP2.request/5` failure, returns `{:error, reason, conn}` so the
  caller can still manage the connection.
  """
  @spec open(Mint.HTTP2.t(), String.t(), String.t(), keyword()) ::
          {:ok, t()} | {:error, term()} | {:error, term(), Mint.HTTP2.t()}
  def open(conn, basin, stream, opts \\ []) do
    Logger.debug(
      "S2S.ReadSession.open basin=#{basin} stream=#{stream} opts=#{inspect(Keyword.delete(opts, :token))}"
    )

    query = Shared.build_read_query(opts)
    path = Shared.records_path(stream) <> query
    token = Keyword.get(opts, :token)
    recv_timeout = Keyword.get(opts, :recv_timeout, Shared.default_timeout())
    headers = Shared.build_headers(basin, token)

    case Mint.HTTP2.request(conn, "GET", path, headers, nil) do
      {:ok, conn, request_ref} ->
        session = %__MODULE__{
          conn: conn,
          request_ref: request_ref,
          owner_pid: self(),
          recv_timeout: recv_timeout
        }

        wait_for_headers(session)

      {:error, conn, reason} ->
        {:error, reason, conn}
    end
  end

  @doc """
  Receive the next batch of records from the session.

  Returns:
  - `{:ok, %ReadBatch{}, session}` — a batch with records
  - `{:error, :end_of_stream, session}` — server closed the stream normally
  - `{:error, :session_closed, session}` — session was already closed
  - `{:error, reason, session}` — an error occurred; session is marked closed
  """
  @spec next_batch(t()) :: {:ok, S2.V1.ReadBatch.t(), t()} | {:error, term(), t()}
  def next_batch(%__MODULE__{closed: true} = session) do
    {:error, :session_closed, session}
  end

  def next_batch(%__MODULE__{} = session) do
    check_owner!(session)

    case Shared.decode_read_batch(session.data) do
      {:ok, batch, rest} ->
        {:ok, batch, %{session | data: rest}}

      :incomplete ->
        receive_batch(session)

      {:error, reason} ->
        {:error, reason, Shared.close_session(session)}
    end
  end

  @doc """
  Close the read session. Sends an HTTP/2 stream cancel (RST_STREAM) to the
  server so it stops sending data. After closing, `next_batch/1` will return
  `{:error, :session_closed, session}`.
  """
  @spec close(t()) :: {:ok, t()}
  def close(%__MODULE__{closed: true} = session), do: {:ok, session}

  def close(%__MODULE__{} = session) do
    conn =
      case Mint.HTTP2.cancel_request(session.conn, session.request_ref) do
        {:ok, conn} ->
          conn

        {:error, conn, reason} ->
          # Best-effort close: stream may already be closed by server.
          Logger.debug("ReadSession.close cancel_request failed: #{inspect(reason)}")
          conn
      end

    {:ok, %{session | conn: conn, closed: true}}
  end

  defp wait_for_headers(session) do
    case Shared.wait_for_headers(session.conn, session.request_ref, session.recv_timeout) do
      {:ok, 200, data, conn} ->
        {:ok, %{session | conn: conn, data: data}}

      {:ok, status, _data, conn} ->
        {:error, %S2.Error{status: status, message: "unexpected status #{status}"}, conn}

      {:error, reason, conn} ->
        {:error, reason, conn}
    end
  end

  defp receive_batch(session) do
    receive_batch(session, Shared.deadline(session.recv_timeout))
  end

  @doc false
  def handle_batch_response({:ok, conn, responses}, session, acc) do
    session = %{session | conn: conn}
    new_data = Shared.extract_data(responses, session.request_ref)
    all_data = acc <> new_data

    case Shared.check_buffer_size(all_data) do
      {:error, :buffer_overflow} ->
        {:error, :buffer_overflow, Shared.close_session(session)}

      :ok ->
        done? = Shared.done?(responses, session.request_ref)

        case Shared.decode_read_batch(all_data) do
          {:ok, batch, rest} ->
            {:ok, batch, %{session | data: rest}}

          :incomplete when done? ->
            {:error, :end_of_stream, Shared.close_session(session)}

          :incomplete ->
            {:continue, %{session | data: all_data}, all_data}

          {:error, reason} ->
            {:error, reason, Shared.close_session(session)}
        end
    end
  end

  def handle_batch_response({:error, conn, _error, _responses}, session, _acc) do
    {:error, :stream_error, Shared.close_session(session, conn)}
  end

  def handle_batch_response(:unknown, session, acc) do
    {:continue, session, acc}
  end

  defp receive_batch(session, dl) do
    receive do
      message ->
        case handle_batch_response(
               Mint.HTTP2.stream(session.conn, message),
               session,
               session.data
             ) do
          {:continue, session, _acc} -> receive_batch(session, dl)
          result -> result
        end
    after
      Shared.remaining(dl) -> {:error, :timeout, Shared.close_session(session)}
    end
  end

  defp check_owner!(%__MODULE__{owner_pid: pid}) do
    Shared.assert_owner!(pid, "ReadSession")
  end
end