codex_sdk

defmodule Codex.StreamQueue do
  @moduledoc false

  use GenServer

  @type t :: pid()

  @spec start_link(keyword()) :: {:ok, pid()} | {:error, term()}
  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, :ok, opts)
  end

  @spec push(t(), term()) :: :ok
  def push(queue, value) when is_pid(queue) do
    GenServer.cast(queue, {:push, value})
  end

  @spec close(t(), term()) :: :ok
  def close(queue, reason \\ :normal) when is_pid(queue) do
    GenServer.cast(queue, {:close, reason})
  end

  @spec pop(t(), timeout()) :: {:ok, term()} | {:error, term()} | :done
  def pop(queue, timeout \\ 5_000) when is_pid(queue) do
    GenServer.call(queue, :pop, timeout)
  end

  @spec stream(t()) :: Enumerable.t()
  def stream(queue) when is_pid(queue) do
    Stream.resource(
      fn -> queue end,
      fn q ->
        case pop(q, :infinity) do
          {:ok, value} -> {[value], q}
          :done -> {:halt, q}
          {:error, reason} when is_exception(reason) -> raise reason
          {:error, reason} -> raise RuntimeError, "stream closed with error: #{inspect(reason)}"
        end
      end,
      fn _ -> :ok end
    )
  end

  @impl true
  def init(:ok) do
    {:ok, %{queue: :queue.new(), closed?: false, waiters: :queue.new(), error: nil}}
  end

  @impl true
  def handle_cast({:push, value}, %{queue: queue, closed?: closed?, waiters: waiters} = state) do
    case :queue.out(waiters) do
      {{:value, from}, remaining} ->
        GenServer.reply(from, {:ok, value})
        {:noreply, %{state | waiters: remaining}}

      {:empty, _} ->
        {:noreply, %{state | queue: :queue.in(value, queue), closed?: closed?}}
    end
  end

  def handle_cast({:close, _reason}, %{closed?: true} = state), do: {:noreply, state}

  def handle_cast({:close, reason}, %{queue: queue, waiters: waiters} = state) do
    error = close_error(reason)
    reply = if error, do: {:error, error}, else: :done

    Enum.each(:queue.to_list(waiters), &GenServer.reply(&1, reply))
    {:noreply, %{state | closed?: true, waiters: :queue.new(), queue: queue, error: error}}
  end

  @impl true
  def handle_call(:pop, _from, %{queue: queue, closed?: true, error: error} = state) do
    reply =
      case error do
        nil -> :done
        reason -> {:error, reason}
      end

    case :queue.out(queue) do
      {{:value, value}, remaining} ->
        {:reply, {:ok, value}, %{state | queue: remaining}}

      {:empty, _} ->
        {:reply, reply, state}
    end
  end

  def handle_call(:pop, from, %{queue: queue, closed?: false, waiters: waiters} = state) do
    case :queue.out(queue) do
      {{:value, value}, remaining} ->
        {:reply, {:ok, value}, %{state | queue: remaining}}

      {:empty, _} ->
        {:noreply, %{state | waiters: :queue.in(from, waiters)}}
    end
  end

  defp close_error({:error, reason}), do: reason
  defp close_error(_), do: nil
end