codex_sdk

defmodule Codex.Realtime.Session do
  @moduledoc """
  Manages a realtime session with WebSocket connection.

  A `RealtimeSession` is a GenServer that manages the connection to OpenAI's
  Realtime API, handles event dispatch, tool execution, and conversation history.

  ## Usage

      # Define an agent
      agent = %Codex.Realtime.Agent{
        name: "Assistant",
        model: "gpt-4o-realtime-preview",
        instructions: "Be helpful and concise."
      }

      # Start the session
      {:ok, session} = Session.start_link(agent: agent)

      # Subscribe to events
      :ok = Session.subscribe(session, self())

      # Send a message
      :ok = Session.send_message(session, "Hello!")

      # Receive events
      receive do
        {:session_event, %Events.AgentStartEvent{}} ->
          IO.puts("Agent started!")

        {:session_event, %Events.AudioEvent{audio: audio}} ->
          # Handle audio data
      end

      # Close when done
      Session.close(session)

  ## Events

  Subscribers receive events as `{:session_event, event}` messages. See
  `Codex.Realtime.Events` for the full list of event types.

  ## Tool Execution

  When the model calls a tool, the session automatically executes it and
  sends the result back to the model. Tool events are emitted to subscribers.
  """

  use GenServer

  alias Codex.Handoff
  alias Codex.Realtime.Agent, as: RealtimeAgent
  alias Codex.Realtime.Config
  alias Codex.Realtime.Config.ModelConfig
  alias Codex.Realtime.Config.SessionModelSettings
  alias Codex.Realtime.Events
  alias Codex.Realtime.Items
  alias Codex.Realtime.ModelEvents
  alias Codex.Realtime.ModelInputs
  alias Codex.Realtime.OpenAIWebSocket
  alias Codex.Realtime.PlaybackTracker

  require Logger

  defstruct [
    :agent,
    :websocket_pid,
    :websocket_module,
    :config,
    :run_config,
    :context,
    :playback_tracker,
    history: [],
    subscribers: %{},
    pending_tool_calls: %{},
    active_response?: false,
    pending_response_create?: false,
    item_transcripts: %{},
    item_guardrail_run_counts: %{},
    interrupted_response_ids: MapSet.new()
  ]

  @type t :: %__MODULE__{
          agent: term(),
          websocket_pid: pid() | nil,
          websocket_module: module(),
          config: ModelConfig.t(),
          run_config: Config.RunConfig.t(),
          context: map(),
          playback_tracker: PlaybackTracker.t(),
          history: [Items.item()],
          subscribers: %{optional(pid()) => reference()},
          pending_tool_calls: %{optional(pid()) => map()},
          active_response?: boolean(),
          pending_response_create?: boolean(),
          item_transcripts: %{String.t() => String.t()},
          item_guardrail_run_counts: %{String.t() => non_neg_integer()},
          interrupted_response_ids: MapSet.t(String.t())
        }

  # Client API

  @doc """
  Start a realtime session.

  ## Options

    * `:agent` - Required. The realtime agent configuration.
    * `:config` - Optional model configuration (API key, URL, etc.).
    * `:run_config` - Optional runtime configuration.
    * `:context` - Optional context map passed to tools and events.
    * `:websocket_pid` - Optional. For testing with a mock WebSocket.
    * `:websocket_module` - Optional. Override the WebSocket transport module
      (default: `Codex.Realtime.OpenAIWebSocket`). Custom modules must provide
      `send_frame/2` and `close/1`.

  ## Returns

    * `{:ok, pid}` on success
    * `{:error, reason}` on failure
  """
  @spec start_link(keyword()) :: GenServer.on_start()
  def start_link(opts) do
    GenServer.start_link(__MODULE__, opts)
  end

  @doc """
  Subscribe to session events.

  The subscriber process will receive `{:session_event, event}` messages
  for all session events.
  """
  @spec subscribe(GenServer.server(), pid()) :: :ok
  def subscribe(session, subscriber_pid) do
    GenServer.call(session, {:subscribe, subscriber_pid})
  end

  @doc """
  Unsubscribe from session events.
  """
  @spec unsubscribe(GenServer.server(), pid()) :: :ok
  def unsubscribe(session, subscriber_pid) do
    GenServer.call(session, {:unsubscribe, subscriber_pid})
  end

  @doc """
  Send audio data to the model.

  Returns `{:error, :not_connected}` if the websocket is unavailable.

  ## Options

    * `:commit` - Whether to commit the audio buffer (default: false)

  ## Examples

      Session.send_audio(session, audio_bytes)
      Session.send_audio(session, audio_bytes, commit: true)
  """
  @spec send_audio(GenServer.server(), binary(), keyword()) :: :ok | {:error, term()}
  def send_audio(session, audio, opts \\ []) do
    commit = Keyword.get(opts, :commit, false)
    GenServer.call(session, {:send_audio, audio, commit})
  end

  @doc """
  Send a text message to the model.

  Returns `{:error, :not_connected}` if the websocket is unavailable.

  Can be a simple string or a structured message map.

  ## Examples

      Session.send_message(session, "Hello!")

      Session.send_message(session, %{
        "type" => "message",
        "role" => "user",
        "content" => [
          %{"type" => "input_text", "text" => "Hello!"},
          %{"type" => "input_image", "image_url" => "data:image/..."}
        ]
      })
  """
  @spec send_message(GenServer.server(), String.t() | map()) :: :ok | {:error, term()}
  def send_message(session, message) do
    GenServer.call(session, {:send_message, message})
  end

  @doc """
  Interrupt the current response.

  Sends a cancel signal to stop the model from generating more output.

  Returns `{:error, :not_connected}` if the websocket is unavailable.
  """
  @spec interrupt(GenServer.server()) :: :ok | {:error, term()}
  def interrupt(session) do
    GenServer.call(session, :interrupt)
  end

  @doc """
  Send a raw event to the model.

  Use this for advanced scenarios where you need to send custom events.

  Returns `{:error, :not_connected}` if the websocket is unavailable.
  """
  @spec send_raw_event(GenServer.server(), map()) :: :ok | {:error, term()}
  def send_raw_event(session, event) do
    GenServer.call(session, {:send_raw_event, event})
  end

  @doc """
  Update session settings.

  Use this to change model settings mid-session, such as voice or modalities.

  Returns `{:error, :not_connected}` if the websocket is unavailable.
  """
  @spec update_session(GenServer.server(), SessionModelSettings.t()) :: :ok | {:error, term()}
  def update_session(session, settings) do
    GenServer.call(session, {:update_session, settings})
  end

  @doc """
  Get the conversation history.
  """
  @spec history(GenServer.server()) :: [Items.item()]
  def history(session) do
    GenServer.call(session, :history)
  end

  @doc """
  Get the current agent.
  """
  @spec current_agent(GenServer.server()) :: term()
  def current_agent(session) do
    GenServer.call(session, :current_agent)
  end

  @doc """
  Close the session.

  Closes the WebSocket connection and stops the session process.
  """
  @spec close(GenServer.server()) :: :ok
  def close(session) do
    GenServer.stop(session, :normal)
  end

  # GenServer Callbacks

  @impl true
  def init(opts) do
    Process.flag(:trap_exit, true)

    agent = Keyword.fetch!(opts, :agent)
    config = Keyword.get(opts, :config, %ModelConfig{})
    run_config = Keyword.get(opts, :run_config, %Config.RunConfig{})
    context = Keyword.get(opts, :context, %{})

    # Support for testing with mock WebSocket
    websocket_pid = Keyword.get(opts, :websocket_pid)
    websocket_module = Keyword.get(opts, :websocket_module) || OpenAIWebSocket

    state = %__MODULE__{
      agent: agent,
      config: config,
      run_config: run_config,
      context: context,
      playback_tracker: PlaybackTracker.new(),
      websocket_pid: websocket_pid,
      websocket_module: websocket_module
    }

    if websocket_pid do
      {:ok, state}
    else
      {:ok, state, {:continue, :connect_websocket}}
    end
  end

  @impl true
  def handle_continue(:connect_websocket, state) do
    case start_websocket(state) do
      {:ok, ws_pid} ->
        {:noreply, %{state | websocket_pid: ws_pid}}

      {:error, reason} ->
        {:stop, reason, state}
    end
  end

  @impl true
  def handle_call({:subscribe, pid}, _from, state) do
    case Map.fetch(state.subscribers, pid) do
      {:ok, _ref} ->
        {:reply, :ok, state}

      :error ->
        ref = Process.monitor(pid)
        {:reply, :ok, %{state | subscribers: Map.put(state.subscribers, pid, ref)}}
    end
  end

  def handle_call({:unsubscribe, pid}, _from, state) do
    case Map.pop(state.subscribers, pid) do
      {nil, _subscribers} ->
        {:reply, :ok, state}

      {ref, subscribers} ->
        Process.demonitor(ref, [:flush])
        {:reply, :ok, %{state | subscribers: subscribers}}
    end
  end

  def handle_call({:send_audio, audio, commit}, _from, state) do
    case send_to_websocket(state, ModelInputs.send_audio(audio, commit)) do
      :ok -> {:reply, :ok, state}
      {:error, _} = error -> {:reply, error, state}
    end
  end

  def handle_call({:send_message, message}, _from, state) do
    with :ok <- send_to_websocket(state, ModelInputs.send_user_input(message)),
         {:ok, state} <- request_response_create(state) do
      {:reply, :ok, state}
    else
      {:error, reason, state} ->
        {:reply, {:error, reason}, state}

      {:error, _} = error ->
        {:reply, error, state}
    end
  end

  def handle_call(:interrupt, _from, state) do
    case send_to_websocket(state, ModelInputs.send_interrupt()) do
      :ok ->
        state = %{
          state
          | playback_tracker: PlaybackTracker.on_interrupted(state.playback_tracker)
        }

        {:reply, :ok, state}

      {:error, _} = error ->
        {:reply, error, state}
    end
  end

  def handle_call({:send_raw_event, event}, _from, state) do
    case send_to_websocket(state, ModelInputs.send_raw_message(event)) do
      :ok -> {:reply, :ok, track_outbound_event(state, event)}
      {:error, _} = error -> {:reply, error, state}
    end
  end

  def handle_call({:update_session, settings}, _from, state) do
    case send_to_websocket(state, ModelInputs.send_session_update(settings)) do
      :ok -> {:reply, :ok, state}
      {:error, _} = error -> {:reply, error, state}
    end
  end

  def handle_call(:history, _from, state) do
    {:reply, state.history, state}
  end

  def handle_call(:current_agent, _from, state) do
    {:reply, state.agent, state}
  end

  @impl true
  def handle_info({:model_event, event}, state) do
    state = handle_model_event(event, state)
    {:noreply, state}
  end

  def handle_info({:websocket_event, json}, state) do
    state = maybe_emit_response_done_error(json, state)

    case ModelEvents.from_json(json) do
      {:ok, event} ->
        state = handle_model_event(event, state)
        {:noreply, state}

      {:error, _} ->
        {:noreply, state}
    end
  end

  def handle_info({:EXIT, pid, reason}, %{websocket_pid: pid} = state) do
    event =
      Events.error(
        %{"type" => "websocket_exit", "reason" => format_reason(reason)},
        state.context
      )

    notify_subscribers(state, event)
    {:noreply, %{state | websocket_pid: nil}}
  end

  def handle_info({:EXIT, _pid, _reason}, state) do
    {:noreply, state}
  end

  def handle_info({:tool_call_result, pid, output}, state) when is_pid(pid) do
    case Map.pop(state.pending_tool_calls, pid) do
      {nil, _pending} ->
        {:noreply, state}

      {pending_tool_call, pending} ->
        Process.demonitor(pending_tool_call.monitor_ref, [:flush])
        state = %{state | pending_tool_calls: pending}
        state = finish_tool_call(state, pending_tool_call, output)
        {:noreply, state}
    end
  end

  def handle_info({:DOWN, ref, :process, pid, reason}, state) do
    case pop_subscriber_by_ref(state.subscribers, pid, ref) do
      {:ok, subscribers} ->
        {:noreply, %{state | subscribers: subscribers}}

      :error ->
        handle_tool_call_down(state, ref, pid, reason)
    end
  end

  def handle_info(_msg, state) do
    {:noreply, state}
  end

  @impl true
  def terminate(_reason, state) do
    _ = drain_pending_tool_calls(state)
    close_websocket(state.websocket_module, state.websocket_pid)

    :ok
  end

  # Private Functions

  defp start_websocket(state) do
    model_name = get_model_name(state.agent)

    OpenAIWebSocket.start_link(
      session_pid: self(),
      config: state.config,
      model_name: model_name
    )
  end

  defp get_model_name(%{model: model}) when is_binary(model), do: model
  defp get_model_name(_), do: RealtimeAgent.default_model()

  defp send_to_websocket(state, input) do
    if state.websocket_pid do
      json = ModelInputs.to_json(input)
      do_send_to_websocket(state.websocket_module, state.websocket_pid, json)
    else
      {:error, :not_connected}
    end
  end

  defp do_send_to_websocket(ws_module, pid, msgs) when is_list(msgs) do
    Enum.reduce_while(msgs, :ok, fn msg, :ok ->
      case do_send_to_websocket(ws_module, pid, msg) do
        :ok -> {:cont, :ok}
        {:error, _} = error -> {:halt, error}
      end
    end)
  end

  defp do_send_to_websocket(ws_module, pid, msg) when is_map(msg) do
    case ws_module.send_frame(pid, {:text, Jason.encode!(msg)}) do
      :ok ->
        :ok

      {:error, _} = error ->
        error

      other ->
        {:error, {:send_failed, other}}
    end
  catch
    :exit, _ -> {:error, :not_connected}
  end

  defp close_websocket(_ws_module, nil), do: :ok

  defp close_websocket(ws_module, pid) do
    ws_module.close(pid)
  catch
    :exit, _ -> :ok
  end

  defp maybe_emit_response_done_error(
         %{"type" => "response.done", "response" => %{"status" => "failed"} = response},
         state
       ) do
    error = response_done_error(response)
    notify_subscribers(state, Events.error(error, state.context))
    state
  end

  defp maybe_emit_response_done_error(_json, state), do: state

  defp response_done_error(%{} = response) do
    base =
      case Map.get(response, "status_details") do
        %{"error" => %{} = error} -> error
        %{"error" => error} when is_binary(error) -> %{"message" => error}
        %{"reason" => reason} when is_binary(reason) -> %{"message" => reason}
        _ -> %{}
      end

    base
    |> maybe_put_error_field("status", response["status"])
    |> maybe_put_error_field("response_id", response["id"])
    |> Map.put("source_event", "response.done")
    |> ensure_error_message(response)
  end

  defp maybe_put_error_field(map, _key, nil), do: map
  defp maybe_put_error_field(map, key, value), do: Map.put(map, key, value)

  defp ensure_error_message(%{"message" => message} = error, _response)
       when is_binary(message) and byte_size(message) > 0 do
    error
  end

  defp ensure_error_message(error, response) do
    Map.put(error, "message", "response.done failed (status=#{inspect(response["status"])})")
  end

  # Event Handlers

  defp handle_model_event(%ModelEvents.ConnectionStatusEvent{status: :connected}, state) do
    # Send initial session configuration
    send_initial_config(state)
    state
  end

  defp handle_model_event(%ModelEvents.ItemUpdatedEvent{item: item}, state) do
    # Check if this is a new item or an update
    is_new = not Enum.any?(state.history, &(&1.item_id == item.item_id))

    # Update history
    history = update_history(state.history, item)

    if is_new do
      event = Events.history_added(item, state.context)
      notify_subscribers(state, event)
    else
      event = Events.history_updated(history, state.context)
      notify_subscribers(state, event)
    end

    %{state | history: history}
  end

  defp handle_model_event(%ModelEvents.ItemDeletedEvent{item_id: item_id}, state) do
    history = Enum.reject(state.history, &(&1.item_id == item_id))
    event = Events.history_updated(history, state.context)
    notify_subscribers(state, event)
    %{state | history: history}
  end

  defp handle_model_event(%ModelEvents.AudioEvent{} = audio, state) do
    event = Events.audio(audio, audio.item_id, audio.content_index, state.context)
    notify_subscribers(state, event)
    state
  end

  defp handle_model_event(%ModelEvents.AudioDoneEvent{} = done, state) do
    event = Events.audio_end(done.item_id, done.content_index, state.context)
    notify_subscribers(state, event)
    state
  end

  defp handle_model_event(%ModelEvents.AudioInterruptedEvent{} = interrupted, state) do
    event =
      Events.audio_interrupted(interrupted.item_id, interrupted.content_index, state.context)

    notify_subscribers(state, event)
    %{state | playback_tracker: PlaybackTracker.on_interrupted(state.playback_tracker)}
  end

  defp handle_model_event(%ModelEvents.ToolCallEvent{} = tool_call, state) do
    execute_tool_call(tool_call, state)
  end

  defp handle_model_event(%ModelEvents.TranscriptDeltaEvent{} = delta, state) do
    # Accumulate transcript for guardrail debouncing
    item_id = delta.item_id
    current_transcript = Map.get(state.item_transcripts, item_id, "")
    new_transcript = current_transcript <> delta.delta

    item_transcripts = Map.put(state.item_transcripts, item_id, new_transcript)
    state = %{state | item_transcripts: item_transcripts}

    # Update history with transcript
    content = [Items.assistant_audio(nil, new_transcript)]

    item =
      Items.assistant_message(item_id, content, status: :in_progress)

    history = update_history(state.history, item)
    %{state | history: history}
  end

  defp handle_model_event(%ModelEvents.TurnStartedEvent{}, state) do
    event = Events.agent_start(state.agent, state.context)
    notify_subscribers(state, event)
    %{state | active_response?: true}
  end

  defp handle_model_event(%ModelEvents.TurnEndedEvent{}, state) do
    # Clear guardrail state for next turn
    event = Events.agent_end(state.agent, state.context)
    notify_subscribers(state, event)

    state = %{
      state
      | active_response?: false,
        item_transcripts: %{},
        item_guardrail_run_counts: %{}
    }

    maybe_flush_deferred_response_create(state)
  end

  defp handle_model_event(%ModelEvents.ErrorEvent{error: error}, state) do
    event = Events.error(error, state.context)
    notify_subscribers(state, event)
    state
  end

  defp handle_model_event(%ModelEvents.InputAudioTranscriptionCompletedEvent{} = tc, state) do
    # Update history with completed transcription
    history = update_history_with_transcription(state.history, tc.item_id, tc.transcript)
    %{state | history: history}
  end

  defp handle_model_event(%ModelEvents.InputAudioTimeoutTriggeredEvent{}, state) do
    event = Events.input_audio_timeout_triggered(state.context)
    notify_subscribers(state, event)
    state
  end

  defp handle_model_event(
         %ModelEvents.RawServerEvent{
           data: %{"type" => "response.done", "response" => %{"status" => "failed"} = response}
         } = event,
         state
       ) do
    notify_subscribers(state, Events.error(response_done_error(response), state.context))

    wrapped = Events.raw_model_event(event, state.context)
    notify_subscribers(state, wrapped)
    state
  end

  defp handle_model_event(event, state) do
    # Wrap other events as raw model events
    wrapped = Events.raw_model_event(event, state.context)
    notify_subscribers(state, wrapped)
    state
  end

  defp send_initial_config(state) do
    # Resolve instructions
    instructions = resolve_instructions(state.agent, state.context)

    settings = %SessionModelSettings{
      model_name: get_model_name(state.agent),
      instructions: instructions,
      tools: get_agent_tools(state.agent),
      modalities: [:text, :audio],
      input_audio_format: :pcm16,
      output_audio_format: :pcm16
    }

    # Merge with run_config settings
    settings =
      case state.run_config.model_settings do
        nil -> settings
        override -> Config.merge_settings(settings, override)
      end

    send_to_websocket(state, ModelInputs.send_session_update(settings))
  end

  defp resolve_instructions(%{instructions: instructions}, _context)
       when is_binary(instructions) do
    instructions
  end

  defp resolve_instructions(%{instructions: instructions}, context)
       when is_function(instructions, 1) do
    instructions.(context)
  end

  defp resolve_instructions(_, _), do: ""

  defp get_agent_tools(%{tools: tools, handoffs: handoffs})
       when is_list(tools) and is_list(handoffs) do
    tools ++ handoff_tools(handoffs)
  end

  defp get_agent_tools(%{tools: tools}) when is_list(tools), do: tools
  defp get_agent_tools(%{handoffs: handoffs}) when is_list(handoffs), do: handoff_tools(handoffs)
  defp get_agent_tools(_), do: []

  defp handoff_tools(handoffs) when is_list(handoffs) do
    handoffs
    |> Enum.map(&handoff_tool_schema/1)
    |> Enum.reject(&is_nil/1)
  end

  defp handoff_tool_schema(handoff) do
    case handoff_tool_name(handoff) do
      name when is_binary(name) ->
        %{
          "type" => "function",
          "name" => name,
          "description" => handoff_tool_description(handoff),
          "parameters" => handoff_input_schema(handoff)
        }

      _ ->
        nil
    end
  end

  defp handoff_tool_name(%Handoff{tool_name: name}) when is_binary(name), do: name

  defp handoff_tool_name(%RealtimeAgent{name: name}) when is_binary(name),
    do: default_handoff_tool_name(name)

  defp handoff_tool_name(%{name: name}) when is_binary(name), do: default_handoff_tool_name(name)

  defp handoff_tool_name(%{"name" => name}) when is_binary(name),
    do: default_handoff_tool_name(name)

  defp handoff_tool_name(_), do: nil

  defp handoff_tool_description(%Handoff{tool_description: description})
       when is_binary(description),
       do: description

  defp handoff_tool_description(%Handoff{agent_name: agent_name}),
    do: default_handoff_tool_description(agent_name, nil)

  defp handoff_tool_description(%RealtimeAgent{name: name, handoff_description: description}),
    do: default_handoff_tool_description(name, description)

  defp handoff_tool_description(%{name: name, handoff_description: description}),
    do: default_handoff_tool_description(name, description)

  defp handoff_tool_description(%{"name" => name, "handoff_description" => description}),
    do: default_handoff_tool_description(name, description)

  defp handoff_tool_description(%{name: name}),
    do: default_handoff_tool_description(name, nil)

  defp handoff_tool_description(%{"name" => name}),
    do: default_handoff_tool_description(name, nil)

  defp handoff_tool_description(_), do: "Handoff to another agent."

  defp handoff_input_schema(%Handoff{input_schema: %{} = schema}) when map_size(schema) > 0,
    do: schema

  defp handoff_input_schema(_), do: %{"type" => "object", "properties" => %{}}

  defp default_handoff_tool_name(name) do
    safe_name =
      name
      |> to_string()
      |> String.trim()
      |> String.replace(~r/\s+/, "_")
      |> String.replace(~r/[^a-zA-Z0-9_]/, "_")
      |> String.downcase()

    "transfer_to_#{safe_name}"
  end

  defp default_handoff_tool_description(name, description) do
    agent_name =
      case name do
        value when is_binary(value) -> value
        value when is_atom(value) -> Atom.to_string(value)
        value -> inspect(value)
      end

    base = "Handoff to the #{agent_name} agent to handle the request."

    description_text =
      case description do
        nil -> ""
        value when is_binary(value) -> String.trim(value)
        value -> value |> inspect() |> String.trim()
      end

    if description_text == "" do
      base
    else
      "#{base} #{description_text}"
    end
  end

  defp update_history(history, item) do
    case Enum.find_index(history, &(&1.item_id == item.item_id)) do
      nil -> history ++ [item]
      idx -> List.replace_at(history, idx, item)
    end
  end

  defp update_history_with_transcription(history, item_id, transcript) do
    Enum.map(history, fn item ->
      if item.item_id == item_id do
        update_item_transcript(item, transcript)
      else
        item
      end
    end)
  end

  defp update_item_transcript(%Items.UserMessageItem{content: content} = item, transcript) do
    updated_content =
      Enum.map(content, fn
        %Items.InputAudio{} = audio -> %{audio | transcript: transcript}
        other -> other
      end)

    %{item | content: updated_content}
  end

  defp update_item_transcript(item, _transcript), do: item

  defp execute_tool_call(tool_call, state) do
    tools = get_agent_tools(state.agent)
    tool = find_tool(tools, tool_call.name)

    event = Events.tool_start(state.agent, tool, tool_call.arguments, state.context)
    notify_subscribers(state, event)

    case handle_handoff_tool_call(tool_call, tool, state) do
      {:ok, state, output} ->
        finish_tool_call(state, %{tool_call: tool_call, tool: tool}, output)

      :not_handoff ->
        {:ok, pid} =
          start_tool_task(fn ->
            resolve_tool_output(tool, tool_call, state.context)
          end)

        monitor_ref = Process.monitor(pid)

        pending_tool_call = %{
          pid: pid,
          monitor_ref: monitor_ref,
          tool_call: tool_call,
          tool: tool
        }

        %{state | pending_tool_calls: Map.put(state.pending_tool_calls, pid, pending_tool_call)}
    end
  end

  defp handle_handoff_tool_call(tool_call, _tool, state) do
    case resolve_handoff_target(state.agent, tool_call.name, tool_call.arguments, state.context) do
      {:ok, nil} ->
        {:ok, state, "Error: Handoff failed - target agent was nil"}

      {:ok, to_agent} ->
        from_agent = state.agent
        state = %{state | agent: to_agent}

        send_initial_config(state)
        notify_subscribers(state, Events.handoff(from_agent, to_agent, state.context))

        output = "Handoff complete. You are now connected to #{agent_display_name(to_agent)}."
        {:ok, state, output}

      {:error, :not_found} ->
        :not_handoff

      {:error, reason} ->
        {:ok, state, "Error: Handoff failed - #{format_reason(reason)}"}
    end
  end

  defp resolve_handoff_target(agent, tool_name, arguments_json, context) do
    case find_handoff(get_agent_handoffs(agent), tool_name) do
      nil ->
        {:error, :not_found}

      %Handoff{agent: target_agent} when not is_nil(target_agent) ->
        {:ok, target_agent}

      %Handoff{on_invoke_handoff: invoke} when is_function(invoke, 2) ->
        args = decode_tool_arguments(arguments_json)
        {:ok, invoke.(context, args)}

      %Handoff{on_invoke_handoff: invoke} when is_function(invoke, 1) ->
        {:ok, invoke.(context)}

      target ->
        {:ok, target}
    end
  rescue
    error -> {:error, error}
  end

  defp get_agent_handoffs(%{handoffs: handoffs}) when is_list(handoffs), do: handoffs
  defp get_agent_handoffs(_), do: []

  defp find_handoff(handoffs, tool_name) do
    Enum.find(handoffs, fn handoff ->
      handoff_tool_name(handoff) == tool_name
    end)
  end

  defp decode_tool_arguments(arguments_json) when is_binary(arguments_json) do
    case Jason.decode(arguments_json) do
      {:ok, args} when is_map(args) -> args
      {:ok, _} -> %{}
      {:error, _} -> %{}
    end
  end

  defp decode_tool_arguments(_), do: %{}

  defp agent_display_name(%{name: name}) when is_binary(name) and name != "", do: name
  defp agent_display_name(%{"name" => name}) when is_binary(name) and name != "", do: name
  defp agent_display_name(_), do: "the selected agent"

  defp resolve_tool_output(nil, tool_call, _context) do
    "Error: Unknown tool #{tool_call.name}"
  end

  defp resolve_tool_output(tool, tool_call, context) do
    execute_tool(tool, tool_call.arguments, context)
  end

  defp finish_tool_call(state, pending_tool_call, output) do
    event =
      Events.tool_end(
        state.agent,
        pending_tool_call.tool,
        pending_tool_call.tool_call.arguments,
        output,
        state.context
      )

    notify_subscribers(state, event)

    state
    |> maybe_send_tool_output(pending_tool_call.tool_call, output)
    |> maybe_start_response_after_tool_output()
  end

  defp maybe_send_tool_output(state, tool_call, output) do
    case send_to_websocket(state, ModelInputs.send_tool_output(tool_call, output, false)) do
      :ok -> state
      {:error, _} -> state
    end
  end

  defp maybe_start_response_after_tool_output(state) do
    case request_response_create(state) do
      {:ok, state} -> state
      {:error, _reason, state} -> state
    end
  end

  defp maybe_flush_deferred_response_create(%{pending_response_create?: false} = state), do: state

  defp maybe_flush_deferred_response_create(state) do
    state
    |> Map.put(:pending_response_create?, false)
    |> request_response_create()
    |> case do
      {:ok, state} -> state
      {:error, _reason, state} -> state
    end
  end

  defp request_response_create(%{active_response?: true} = state) do
    {:ok, %{state | pending_response_create?: true}}
  end

  defp request_response_create(state) do
    case send_response_create(state) do
      :ok ->
        {:ok, %{state | active_response?: true, pending_response_create?: false}}

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

  defp send_response_create(state) do
    send_to_websocket(state, ModelInputs.send_raw_message(%{"type" => "response.create"}))
  end

  defp track_outbound_event(state, %{"type" => "response.create"}) do
    %{state | active_response?: true, pending_response_create?: false}
  end

  defp track_outbound_event(state, _event), do: state

  defp handle_tool_call_down(state, ref, pid, :normal) do
    case Map.pop(state.pending_tool_calls, pid) do
      {nil, _pending} ->
        {:noreply, state}

      {%{monitor_ref: ^ref}, pending} ->
        {:noreply, %{state | pending_tool_calls: pending}}

      {pending_tool_call, pending} ->
        Process.demonitor(pending_tool_call.monitor_ref, [:flush])
        {:noreply, %{state | pending_tool_calls: pending}}
    end
  end

  defp handle_tool_call_down(state, ref, pid, reason) do
    case Map.pop(state.pending_tool_calls, pid) do
      {nil, _pending} ->
        {:noreply, state}

      {%{monitor_ref: ^ref} = pending_tool_call, pending} ->
        state = %{state | pending_tool_calls: pending}
        output = "Error: Tool execution failed - #{format_reason(reason)}"
        state = finish_tool_call(state, pending_tool_call, output)
        {:noreply, state}

      {pending_tool_call, pending} ->
        Process.demonitor(pending_tool_call.monitor_ref, [:flush])
        state = %{state | pending_tool_calls: pending}
        output = "Error: Tool execution failed - #{format_reason(reason)}"
        state = finish_tool_call(state, pending_tool_call, output)
        {:noreply, state}
    end
  end

  defp drain_pending_tool_calls(state) do
    Enum.each(state.pending_tool_calls, fn
      {pid, %{monitor_ref: monitor_ref}} ->
        Process.demonitor(monitor_ref, [:flush])

        if is_pid(pid) and Process.alive?(pid) do
          Process.exit(pid, :shutdown)
        end

      {_pid, _pending_tool_call} ->
        :ok
    end)

    :ok
  end

  @spec start_tool_task((-> any())) :: {:ok, pid()}
  defp start_tool_task(fun) do
    parent = self()

    runner = fn ->
      output = fun.()
      send(parent, {:tool_call_result, self(), output})
    end

    try do
      case Task.Supervisor.start_child(Codex.TaskSupervisor, runner) do
        {:ok, pid} -> {:ok, pid}
        {:error, {:already_started, pid}} -> {:ok, pid}
        {:error, _} -> Task.start_link(runner)
      end
    catch
      :exit, _ -> Task.start_link(runner)
    end
  end

  defp find_tool(tools, name) do
    Enum.find(tools, fn tool ->
      case tool do
        %{name: ^name} -> true
        %{"name" => ^name} -> true
        %Handoff{tool_name: ^name} -> true
        _ -> false
      end
    end)
  end

  defp execute_tool(tool, arguments_json, context) do
    case Jason.decode(arguments_json) do
      {:ok, args} ->
        try do
          result = invoke_tool(tool, args, context)
          to_string(result)
        rescue
          e ->
            "Error: #{Exception.message(e)}"
        end

      {:error, reason} ->
        "Error: Invalid JSON arguments - #{inspect(reason)}"
    end
  end

  defp invoke_tool(%{on_invoke: fun}, args, context) when is_function(fun, 2) do
    fun.(args, context)
  end

  defp invoke_tool(%{on_invoke: fun}, args, _context) when is_function(fun, 1) do
    fun.(args)
  end

  defp invoke_tool(%{execute: fun}, args, context) when is_function(fun, 2) do
    fun.(args, context)
  end

  defp invoke_tool(%{execute: fun}, args, _context) when is_function(fun, 1) do
    fun.(args)
  end

  defp invoke_tool(%{handler: fun}, args, context) when is_function(fun, 2) do
    fun.(args, context)
  end

  defp invoke_tool(%{handler: fun}, args, _context) when is_function(fun, 1) do
    fun.(args)
  end

  defp invoke_tool(_tool, _args, _context) do
    "Error: Tool has no invokable function"
  end

  defp notify_subscribers(state, event) do
    Enum.each(Map.keys(state.subscribers), fn pid ->
      send(pid, {:session_event, event})
    end)
  end

  defp pop_subscriber_by_ref(subscribers, pid, ref) do
    case Map.get(subscribers, pid) do
      ^ref -> {:ok, Map.delete(subscribers, pid)}
      _ -> :error
    end
  end

  defp format_reason(reason) when is_atom(reason), do: Atom.to_string(reason)
  defp format_reason(reason), do: inspect(reason)
end