fnord

defmodule AI.Completion do
  @moduledoc """
  This module sends a request to the model and handles the response. It is able
  to handle tool calls and responses.

  ## Input options

  - `toolbox` - a map of tool names to modules implementing `AI.Tools`; the specs list
    is derived automatically via `AI.Tools.toolbox_to_specs/1`.

  ## Output options

  Output is controlled by the following mechanisms.

  1. `log_msgs` - log messages from the user and assistant as `info`
  2. `log_tool_calls` - log tool calls as `info` and tool call results as `debug`

  `LOGGER_LEVEL` must be set to `debug` to see the output of tool call results.
  """
  import AI.Util

  defstruct [
    :model,
    :web_search?,
    :response_format,
    :toolbox,
    :specs,
    :log_msgs,
    :log_tool_calls,
    :archive_notes,
    :replay_conversation,
    :name,
    :verbosity,
    :usage,
    :messages,
    :tool_call_requests,
    :response,
    :compact?,
    :is_compacting?,
    :conversation_pid,
    # Snapshot of length(messages) at the start of this completion round.
    # tools_used/1 uses this as a stable round delimiter so mid-loop interrupt
    # injection (which appends user messages into state.messages between LLM
    # iterations) does not hide tool calls that fired earlier in the same
    # round. See tools_used/1 for the rationale.
    :initial_message_count,
    # Tool-round counter + cap. One round == one `{:ok, :tool, _}` response
    # from the model. When the counter hits the cap, the next send_request
    # drops the tool surface and injects a system nudge telling the model
    # to produce its final answer. Prevents runaway verify-loops in
    # sub-agents that keep deciding to "check one more thing" without
    # emitting a final response. See maybe_nudge_wrap_up/1.
    :tool_round_count,
    :tool_round_cap
  ]

  @type t :: %__MODULE__{
          model: String.t(),
          web_search?: boolean,
          response_format: map | nil,
          toolbox: AI.Tools.toolbox() | nil,
          specs: list(AI.Tools.tool_spec()) | nil,
          log_msgs: boolean(),
          log_tool_calls: boolean(),
          archive_notes: boolean(),
          replay_conversation: boolean(),
          name: String.t() | nil,
          verbosity: String.t() | nil,
          usage: integer(),
          messages: list(AI.Util.msg()),
          tool_call_requests: list(),
          response: String.t() | nil,
          compact?: bool,
          is_compacting?: bool,
          initial_message_count: non_neg_integer(),
          tool_round_count: non_neg_integer(),
          tool_round_cap: pos_integer()
        }

  @type response ::
          {:ok, t}
          | {:error, t}
          | {:error, binary}
          | {:error, :context_length_exceeded, non_neg_integer}

  @tool_output_preview 1024

  # Cap on tool-call rounds per get/1 invocation. One round is one
  # `{:ok, :tool, _}` response that triggers a recursive `send_request`.
  # Default is conservative enough to catch verify-thrashing (observed at
  # 300+ rounds in a runaway coder session) while leaving headroom for
  # legitimately multi-file changes. Override via FNORD_TOOL_ROUND_CAP.
  @default_tool_round_cap 75

  @spec get(Keyword.t()) :: response
  def get(opts) do
    with {:ok, state} <- new(opts) do
      # Note: we do not check the "name" back in. It is associated with this
      # process' pid. If the agent calls `AI.Completion.get/1` again, it will
      # get the same name, maintaining continuity between multiple completion
      # steps.
      state
      |> AI.Completion.Output.replay_conversation()
      |> send_request()
    end
  end

  @spec new(Keyword.t()) :: {:ok, t} | {:error, any}
  def new(opts) do
    with {:ok, model} <- Keyword.fetch(opts, :model),
         {:ok, messages} <- Keyword.fetch(opts, :messages) do
      response_format = Keyword.get(opts, :response_format, nil)
      name = Keyword.get(opts, :name, nil)
      verbosity = Keyword.get(opts, :verbosity, nil)
      compact? = Keyword.get(opts, :compact?, true)
      web_search? = Keyword.get(opts, :web_search?, false)

      toolbox_opt = Keyword.get(opts, :toolbox, nil)

      toolbox =
        cond do
          is_nil(toolbox_opt) -> nil
          is_map(toolbox_opt) && map_size(toolbox_opt) == 0 -> nil
          true -> AI.Tools.build_toolbox(toolbox_opt)
        end

      specs =
        if is_nil(toolbox) do
          nil
        else
          toolbox
          |> Map.values()
          |> Enum.map(& &1.spec())
        end

      log_msgs = Keyword.get(opts, :log_msgs, false)
      replay = Keyword.get(opts, :replay_conversation, true)

      quiet? = Services.Globals.get_env(:fnord, :quiet)
      log_tool_calls = Keyword.get(opts, :log_tool_calls, !quiet?)

      archive? = Keyword.get(opts, :archive_notes, false)
      messages = set_name(messages, name)
      initial_message_count = length(messages)

      # Back-compat: historically this option key was `:conversation` even
      # though the value was a PID. Prefer the explicit `:conversation_pid`
      # going forward.
      conversation_pid =
        Keyword.get(opts, :conversation_pid) ||
          Keyword.get(opts, :conversation)

      state =
        %__MODULE__{
          model: model,
          web_search?: web_search?,
          response_format: response_format,
          toolbox: toolbox,
          specs: specs,
          log_msgs: log_msgs,
          log_tool_calls: log_tool_calls,
          archive_notes: archive?,
          replay_conversation: replay,
          name: name,
          verbosity: verbosity,
          conversation_pid: conversation_pid,
          usage: 0,
          messages: messages,
          tool_call_requests: [],
          response: nil,
          compact?: compact?,
          is_compacting?: false,
          initial_message_count: initial_message_count,
          tool_round_count: 0,
          tool_round_cap: resolve_tool_round_cap()
        }

      {:ok, state}
    end
  end

  @spec new_from_conversation(Store.Project.Conversation.t(), Keyword.t()) ::
          {:ok, t}
          | {:error, :conversation_not_found}
  def new_from_conversation(conversation, opts) do
    if Store.Project.Conversation.exists?(conversation) do
      with {:ok, %{messages: msgs}} <- Store.Project.Conversation.read(conversation) do
        msgs_atom = Util.string_keys_to_atoms(msgs)
        name = agent_name_from_messages(msgs_atom)

        opts
        |> Keyword.put(:messages, msgs_atom)
        |> then(fn o -> if name, do: Keyword.put(o, :name, name), else: o end)
        |> new()
      end
    else
      {:error, :conversation_not_found}
    end
  end

  @doc """
  Returns a map of tool names to the number of times each tool was called
  during this completion round - that is, in messages that were appended after
  AI.Completion.new/1 captured the starting message length.

  This delimiter is stable in the presence of mid-loop interrupt injection. An
  earlier implementation used "tools after the last user message" as the round
  boundary, which broke whenever a user interjection arrived mid-round: the
  injected user message became the new "last user", hiding any tool calls that
  had already executed earlier in the same round. That bug silently dropped
  the editing_tools_used flag and skipped end-of-session worktree merges.

  Falls back to scanning all messages when initial_message_count is missing
  (older state structs that predate this field).
  """
  @spec tools_used(t) :: %{binary => non_neg_integer()}
  def tools_used(%{messages: messages} = state) do
    drop_count = Map.get(state, :initial_message_count) || 0

    messages
    |> Enum.drop(drop_count)
    |> Enum.reduce(%{}, fn
      %{tool_calls: tool_calls}, acc ->
        tool_calls
        |> Enum.reduce(acc, fn %{function: %{name: func}}, acc ->
          Map.update(acc, func, 1, &(&1 + 1))
        end)

      _, acc ->
        acc
    end)
  end

  # -----------------------------------------------------------------------------
  # Completion handling
  # -----------------------------------------------------------------------------
  defp maybe_apply_interrupts(%{conversation_pid: nil} = state), do: state

  defp maybe_apply_interrupts(%{conversation_pid: pid} = state) do
    pid
    |> Services.Conversation.Interrupts.take_all()
    |> case do
      [] ->
        state

      msgs ->
        new_messages = state.messages ++ msgs
        Services.Conversation.replace_msgs(new_messages, pid)

        Enum.each(msgs, fn %{content: msg} ->
          msg
          |> String.replace_prefix("[User Interjection]", "(rude)")
          |> UI.feedback_user()
        end)

        %{state | messages: new_messages}
    end
  end

  @spec send_request(t) :: response
  defp send_request(state) do
    # Inject any pending user interrupts before calling the model
    state = maybe_apply_interrupts(state)

    AI.CompletionAPI.get(
      state.model,
      state.messages,
      state.specs,
      state.response_format,
      state.web_search?,
      state.verbosity
    )
    |> handle_response(state)
  end

  @spec handle_response({:ok, any} | {:error, any}, t) :: response
  defp handle_response({:ok, :msg, response, usage}, state) do
    {:ok,
     %{
       state
       | messages: state.messages ++ [AI.Util.assistant_msg(response)],
         response: response,
         usage: usage,
         is_compacting?: false
     }}
  end

  defp handle_response({:ok, :tool, tool_calls}, state) do
    state
    |> Map.put(:is_compacting?, false)
    |> Map.put(:tool_call_requests, tool_calls)
    |> Map.update(:tool_round_count, 1, &(&1 + 1))
    |> handle_tool_calls()
    |> maybe_apply_interrupts()
    |> maybe_nudge_wrap_up()
    |> send_request()
  end

  defp handle_response({:error, :context_length_exceeded, usage}, %{compact?: false}) do
    {:error, :context_length_exceeded, usage}
  end

  defp handle_response(
         {:error, :context_length_exceeded, usage},
         %{messages: msgs, is_compacting?: false} = state
       ) do
    UI.warn("[compaction] Context length exceeded, compacting conversation and retrying...")

    with {:ok, compacted, new_usage} <- AI.Completion.Compaction.compact(msgs) do
      %{state | messages: compacted, usage: new_usage, is_compacting?: true}
      |> send_request()
    else
      {:error, _reason} -> {:error, :context_length_exceeded, usage}
    end
  end

  defp handle_response({:error, :context_length_exceeded, usage}, _state) do
    {:error, :context_length_exceeded, usage}
  end

  defp handle_response({:error, :api_unavailable, reason}, _state) do
    {:error,
     """
     The OpenAI API is currently unavailable. Please try again later.
     Error message: #{reason}
     """}
  end

  defp handle_response({:error, %{http_status: http_status, code: code, message: msg}}, state) do
    error_msg =
      """
      I encountered an error while processing your request.

      - HTTP Status: #{http_status}
      - Error code: #{code}
      - Message: #{msg}
      """

    {:error, %{state | response: error_msg}}
  end

  defp handle_response({:error, %{http_status: http_status, message: msg}}, state) do
    error_msg =
      """
      I encountered an error while processing your request.

      - HTTP Status: #{http_status}
      - Message: #{msg}
      """

    {:error, %{state | response: error_msg}}
  end

  defp handle_response({:error, %{http_status: http_status, error: msg}}, state)
       when is_binary(msg) do
    error_msg =
      """
      I encountered an error while processing your request.

      - HTTP Status: #{http_status}
      - Error: #{msg}
      """

    {:error, %{state | response: error_msg}}
  end

  defp handle_response({:error, %{http_status: http_status, error: msg}}, state) do
    error_msg =
      """
      I encountered an error while processing your request.

      - HTTP Status: #{http_status}
      - Error: #{inspect(msg, pretty: true)}
      """

    {:error, %{state | response: error_msg}}
  end

  defp handle_response({:error, reason}, state) do
    reason =
      if is_binary(reason) do
        reason
      else
        inspect(reason, pretty: true)
      end

    error_msg =
      """
      I encountered an error while processing your request.

      The error message was:

      #{reason}
      """

    {:error, %{state | response: error_msg}}
  end

  # -----------------------------------------------------------------------------
  # Tool calls
  #
  # Note: We intentionally process and log each tool call as its own
  # assistant/tool message pair, rather than grouping multiple tool calls in a
  # single assistant message as the OpenAI API allows. This makes it easier to
  # guarantee tool results always appear immediately after their request,
  # simplifies auditing, and avoids out-of-order response issues. If you need
  # OpenAI-compatible message grouping, refactor here.
  # -----------------------------------------------------------------------------
  defp handle_tool_calls(%{tool_call_requests: tool_calls} = state) do
    # Deduplicate tool call requests by function name and arguments (canonicalized)
    tool_calls = Enum.uniq_by(tool_calls, &dedupe_key/1)

    {async_calls, serial_calls} =
      Enum.split_with(tool_calls, fn req ->
        AI.Tools.is_async?(req.function.name, state.toolbox)
      end)

    # First handle async tool calls concurrently
    state =
      async_calls
      |> Util.async_stream(
        &handle_tool_call(state, &1),
        ordered: true
      )
      |> collect_tool_call_result_messages(state)

    # Now handle all remaining requests serially
    state =
      serial_calls
      |> Util.async_stream(
        &handle_tool_call(state, &1),
        ordered: true,
        max_concurrency: 1
      )
      |> collect_tool_call_result_messages(state)

    # Clear out the tool call requests and return
    %{state | tool_call_requests: []}
  end

  defp collect_tool_call_result_messages(results, state) do
    messages =
      results
      |> Enum.reduce(state.messages, fn result, acc ->
        case result do
          {:ok, {:ok, req, res}} ->
            acc ++ [req, res]

          {:ok, other} ->
            UI.report_from(
              state.name,
              "Tool call returned unexpected result",
              inspect(other, pretty: true)
            )

            acc

          {:exit, reason} ->
            UI.report_from(
              state.name,
              "Tool call crashed",
              inspect(reason, pretty: true)
            )

            acc

          other ->
            UI.report_from(
              state.name,
              "Tool call produced unknown result",
              inspect(other, pretty: true)
            )

            acc
        end
      end)

    %{state | messages: messages}
  end

  @spec dedupe_key(map()) :: {String.t(), String.t()} | nil
  defp dedupe_key(%{function: %{name: func, arguments: args_json}}) when is_binary(args_json) do
    case SafeJson.decode(args_json) do
      # Re-encode to get consistent ordering
      {:ok, decoded} -> {func, inspect(decoded, custom_options: [sort_maps: true])}
      # Fallback to raw string if not valid JSON
      _ -> {func, args_json}
    end
  end

  defp dedupe_key(_), do: nil

  @spec handle_tool_call(t, AI.Util.tool_call()) :: {
          :ok,
          AI.Util.tool_request_msg(),
          AI.Util.tool_response_msg()
        }
  def handle_tool_call(state, %{id: id, function: %{name: func, arguments: args_json}}) do
    # --------------------------------------------------------------------------
    # Agents' names are associated with their process ID, and tool call
    # requests and results are reported from within the process that performs
    # the tool call. Because `handle_tool_calls` invokes tools within a
    # separate process, we need to associate the agent's name with the process
    # for the logs to display the correct name.
    #
    # If the tool itself invokes a new agent, that agent will be given a new
    # name in `AI.Agent.get_response/1`.
    # --------------------------------------------------------------------------
    Services.NamePool.associate_name(state.name)

    # Now back to your regularly scheduled programming...
    request = AI.Util.assistant_tool_msg(id, func, args_json)

    with {:ok, output} <- perform_tool_call(state, func, args_json) do
      if state.archive_notes do
        Services.Notes.ingest_research(func, args_json, output)
      end

      response = AI.Util.tool_msg(id, func, output)
      {:ok, request, response}
    else
      {:error, reason} ->
        oopsie(state, func, args_json, reason)
        response = AI.Util.tool_msg(id, func, reason)
        {:ok, request, response}

      {:error, :unknown_tool, tool} ->
        oopsie(state, func, args_json, "Invalid tool #{tool}")

        error = """
        Your attempt to call #{func} failed because the tool '#{tool}' was not found.
        Your tool call request supplied the following arguments: #{args_json}.
        Please consult the specifications for your available tools and use only the tools that are listed.
        """

        response = AI.Util.tool_msg(id, func, error)
        {:ok, request, response}

      {:error, :missing_argument, key} ->
        oopsie(state, func, args_json, "Missing required argument #{key}")

        spec =
          with {:ok, spec} <- AI.Tools.tool_spec(func, state.toolbox),
               {:ok, json} <- SafeJson.encode(spec) do
            json
          else
            error -> "Error retrieving specification: #{inspect(error)}"
          end

        error = """
        Your attempt to call #{func} failed because it was missing a required argument, '#{key}'.
        Your tool call request supplied the following arguments: #{args_json}.
        The parameter `#{key}` must be included and cannot be `null` or an empty string.
        The correct specification for the tool call is: #{spec}
        """

        response = AI.Util.tool_msg(id, func, error)
        {:ok, request, response}

      {:error, :invalid_argument, key} ->
        oopsie(state, func, args_json, "Invalid argument #{key}")

        spec =
          with {:ok, spec} <- AI.Tools.tool_spec(func, state.toolbox),
               {:ok, json} <- SafeJson.encode(spec) do
            json
          else
            error -> "Error retrieving specification: #{inspect(error)}"
          end

        error = """
        Your attempt to call #{func} failed because it contained an invalid argument or value for '#{key}'.
        Your tool call request supplied the following arguments: #{args_json}.
        The parameter `#{key}` must be a valid value as specified in the tool's specification.
        The correct specification for the tool call is: #{spec}
        """

        response = AI.Util.tool_msg(id, func, error)
        {:ok, request, response}

      {:error, exit_code, msg} when is_integer(exit_code) ->
        oopsie(state, func, args_json, "External process exited with code #{exit_code}: #{msg}")

        error = """
        Your attempt to call #{func} failed because the external process exited with an error.
        Exit code: #{exit_code}
        Error message: #{msg}
        """

        response = AI.Util.tool_msg(id, func, error)
        {:ok, request, response}
    end
  end

  @spec perform_tool_call(t, binary, binary) :: AI.Tools.tool_result()
  defp perform_tool_call(state, func, args_json) when is_binary(args_json) do
    with {:ok, args} <- SafeJson.decode(args_json) do
      AI.Tools.with_args(
        func,
        args,
        fn args ->
          AI.Completion.Output.on_event(state, :tool_call, {func, args})

          # Execute the tool call and then conditionally offload very large
          # tool outputs to a temp file to avoid keeping gigantic blobs in
          # memory or logs. If offload fails for any reason, we fall back to
          # the original in-memory content.
          result = AI.Tools.perform_tool_call(func, args, state.toolbox)

          result =
            case result do
              {:ok, resp} when is_binary(resp) ->
                {:ok, maybe_offload_tool_output(resp)}

              {:error, reason} when is_binary(reason) ->
                {:error, maybe_offload_tool_output(reason)}

              {:error, code, msg} when is_integer(code) and is_binary(msg) ->
                {:error, code, maybe_offload_tool_output(msg)}

              other ->
                other
            end

          AI.Completion.Output.on_event(state, :tool_call_result, {func, args, result})
          result
        end,
        state.toolbox
      )
    end
  end

  @doc """
  If a tool produced a very large textual output, attempt to write it to a
  temporary file and replace the in-memory content with a short placeholder
  that points to the temp file and includes a preview. Fail silently and return
  the original content on any error.
  """
  def maybe_offload_tool_output(content) when is_binary(content) do
    if String.length(content) <= AI.Util.max_msg_length() do
      content
    else
      preview = content |> :erlang.binary_part(0, @tool_output_preview)

      try do
        tmp = Services.TempFile.mktemp!()

        case File.chmod(tmp, 0o600) do
          :ok -> :ok
          {:error, reason} -> raise "Failed to chmod file #{tmp}: #{inspect(reason)}"
        end

        File.write!(tmp, content)

        "[Large tool output (#{byte_size(content)} bytes) written to #{tmp}. " <>
          "Use file_contents_tool to read the full output.\nPreview:\n" <>
          preview <> "]"
      rescue
        _ ->
          # On any failure while trying to offload, return the original content
          content
      end
    end
  end

  @spec oopsie(t, binary, binary, any) :: any
  defp oopsie(state, tool, args_json, reason) do
    safe_reason =
      if is_binary(reason) do
        reason
      else
        inspect(reason, pretty: true)
      end

    AI.Completion.Output.on_event(
      state,
      :tool_call_error,
      {tool, args_json, {:error, safe_reason}}
    )
  end

  # Updates the system message that identifies the LLM to itself by name and
  # updates it to use the name provided by the `name` arg, if any.
  defp set_name(messages, nil) do
    set_name(messages, Services.NamePool.default_name())
  end

  defp set_name(messages, name) do
    # Normalize keys so role/content are accessible
    messages = Util.string_keys_to_atoms(messages)

    # Check if any system message already sets the name
    has_name =
      Enum.any?(messages, fn msg ->
        if is_system_msg?(msg) do
          case Regex.run(~r/Your name is .+\./, msg.content) do
            nil -> false
            _ -> true
          end
        else
          false
        end
      end)

    if has_name do
      Enum.map(messages, fn msg ->
        if is_system_msg?(msg) do
          case Regex.run(~r/Your name is .+\./, msg.content) do
            nil -> msg
            _ -> AI.Util.system_msg("Your name is #{name}.")
          end
        else
          msg
        end
      end)
    else
      [AI.Util.system_msg("Your name is #{name}.") | messages]
    end
  end

  defp agent_name_from_messages(messages) do
    Enum.find_value(messages, fn
      %{role: role, content: content} when role in ["system", "developer"] ->
        case Regex.run(~r/^Your name is (.*)\.$/, content) do
          [_, name] -> name
          _ -> nil
        end

      _ ->
        nil
    end)
  end

  # Guard against runaway tool loops. Once the model has made
  # `tool_round_cap` rounds of tool calls in a single `get/1`, drop the
  # tool surface (so the model MUST emit text on the next call) and
  # inject a system nudge explaining what just happened. An empty map
  # (NOT nil) is used as the post-cap toolbox so any tool call that
  # still arrives - replayed message, interrupt injection, anomalous
  # model behavior - routes through AI.Tools.tool_module/2's explicit-
  # empty branch and returns `:unknown_tool` instead of falling back to
  # the default basic+MCP toolbox. The nudge's claim that further tool
  # calls will not be executed is thereby enforced in code, not just
  # advertised in the prompt.
  @spec maybe_nudge_wrap_up(t) :: t
  defp maybe_nudge_wrap_up(%__MODULE__{tool_round_count: n, tool_round_cap: cap} = state)
       when n >= cap do
    nudge = """
    You have made #{n} rounds of tool calls in this turn. That is the
    configured cap. Stop calling tools and produce your final response
    now, based on what you have already gathered. If you were verifying,
    re-checking, or re-reading conventions, rely on prior reads in this
    session. Further tool calls will not be executed (the tool surface
    has been withdrawn for the next model call).
    """

    %{
      state
      | messages: state.messages ++ [AI.Util.system_msg(nudge)],
        specs: nil,
        toolbox: %{}
    }
  end

  defp maybe_nudge_wrap_up(state), do: state

  defp resolve_tool_round_cap do
    case Util.Env.get_env("FNORD_TOOL_ROUND_CAP") do
      nil ->
        @default_tool_round_cap

      raw ->
        case Integer.parse(raw) do
          {n, ""} when n > 0 -> n
          _ -> @default_tool_round_cap
        end
    end
  end
end