ragex

defmodule Ragex.Agent.Executor do
  @moduledoc """
  ReAct (Reasoning + Acting) execution loop for agent operations.

  Implements the agent loop:
  1. Build prompt with conversation history and available tools
  2. Call LLM with tools enabled
  3. If response has tool_calls: execute tools, add results, repeat
  4. Return final text response

  ## Usage

      alias Ragex.Agent.{Executor, Memory, ToolSchema}

      # Create session and run agent
      {:ok, session} = Memory.new_session(%{project_path: "/my/project"})
      Memory.add_message(session.id, :system, "You are a code analysis assistant...")
      Memory.add_message(session.id, :user, "Analyze this project for issues")

      {:ok, result} = Executor.run(session.id, [
        max_iterations: 10,
        provider: :deepseek_r1
      ])
  """

  require Logger

  alias Ragex.Agent.{Memory, StreamConsumer, ToolSchema}
  alias Ragex.AI.Config
  alias Ragex.CLI.Colors
  alias Ragex.MCP.Handlers.Tools, as: MCPTools

  import Ragex.MCP.Handlers.Tools, only: [format_reason: 1]

  @default_max_iterations 15
  @default_temperature 0.7
  @default_max_tokens 4096

  @type run_result :: %{
          content: String.t(),
          iterations: non_neg_integer(),
          tool_calls_made: non_neg_integer(),
          usage: map(),
          session_id: String.t()
        }

  @doc """
  Run the agent execution loop.

  ## Parameters

  - `session_id` - Active session ID with conversation history
  - `opts` - Options:
    - `:max_iterations` - Maximum tool call iterations (default: 15)
    - `:provider` - AI provider override (:deepseek_r1, :openai, :anthropic, :ollama)
    - `:tools` - Custom tool list (default: full agent tool set from `ToolSchema`)
    - `:temperature` - LLM temperature (default: 0.7)
    - `:max_tokens` - Max response tokens (default: 4096)
    - `:tool_choice` - Tool selection strategy ("auto", "any", or specific tool name)
    - `:system_prompt_override` - Replace the session's system prompt for this run

  ## Returns

  - `{:ok, result}` - Execution completed with final response
  - `{:error, reason}` - Execution failed
  """
  @spec run(String.t(), keyword()) :: {:ok, run_result()} | {:error, term()}
  def run(session_id, opts \\ []) do
    max_iterations = Keyword.get(opts, :max_iterations, @default_max_iterations)
    provider = get_provider(opts)
    provider_name = get_provider_name(opts)

    tools = get_tools(provider_name, opts)

    state = %{
      session_id: session_id,
      provider: provider,
      provider_name: provider_name,
      tools: tools,
      opts: opts,
      iterations: 0,
      tool_calls_made: 0,
      total_usage: %{prompt_tokens: 0, completion_tokens: 0, total_tokens: 0},
      tool_call_history: %{}
    }

    case execute_loop(state, max_iterations) do
      {:ok, final_state, content} ->
        {:ok,
         %{
           content: content,
           iterations: final_state.iterations,
           tool_calls_made: final_state.tool_calls_made,
           usage: final_state.total_usage,
           session_id: session_id
         }}

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

  @doc """
  Run the agent execution loop with streaming support.

  Same as `run/2` but streams the final AI response in real-time via callbacks.
  Intermediate tool-call steps use blocking `generate` (provider stream parsers
  don't yet handle tool_call deltas), but the final text response is streamed
  chunk-by-chunk so the user gets immediate feedback.

  ## Additional Options

  - `:on_chunk` - `(chunk -> :ok)` callback for real-time content/thinking delivery
  - `:on_phase` - `(:thinking | :answering | :done -> :ok)` phase transition callback
  - `:on_tool_progress` - `(map() -> :ok)` callback when tool calls are being executed

  ## Returns

  Same as `run/2`.
  """
  @spec stream_run(String.t(), keyword()) :: {:ok, run_result()} | {:error, term()}
  def stream_run(session_id, opts \\ []) do
    max_iterations = Keyword.get(opts, :max_iterations, @default_max_iterations)
    provider = get_provider(opts)
    provider_name = get_provider_name(opts)
    tools = get_tools(provider_name, opts)

    state = %{
      session_id: session_id,
      provider: provider,
      provider_name: provider_name,
      tools: tools,
      opts: opts,
      iterations: 0,
      tool_calls_made: 0,
      total_usage: %{prompt_tokens: 0, completion_tokens: 0, total_tokens: 0},
      streaming: true,
      tool_call_history: %{}
    }

    case execute_streaming_loop(state, max_iterations) do
      {:ok, final_state, content} ->
        {:ok,
         %{
           content: content,
           iterations: final_state.iterations,
           tool_calls_made: final_state.tool_calls_made,
           usage: final_state.total_usage,
           session_id: session_id
         }}

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

  @doc """
  Execute a single step of the agent loop.

  Useful for debugging or manual step-through.
  """
  @spec step(String.t(), keyword()) ::
          {:continue, map()} | {:done, String.t()} | {:error, term()}
  def step(session_id, opts \\ []) do
    provider = get_provider(opts)
    provider_name = get_provider_name(opts)
    tools = get_tools(provider_name, opts)

    state = %{
      session_id: session_id,
      provider: provider,
      provider_name: provider_name,
      tools: tools,
      opts: opts,
      iterations: 0,
      tool_calls_made: 0,
      total_usage: %{}
    }

    execute_step(state)
  end

  # Private functions

  # Build Memory.get_context opts, honouring any :context_max_chars passed by caller.
  defp context_opts(state) do
    base = [format: state.provider_name]

    case Keyword.get(state.opts, :context_max_chars) do
      nil -> base
      max -> Keyword.put(base, :max_chars, max)
    end
  end

  # --- Streaming execution loop ---

  defp execute_streaming_loop(state, max_iterations)
       when state.iterations >= max_iterations do
    Logger.warning("Agent reached max iterations (#{max_iterations})")

    case Memory.get_messages(state.session_id, limit: 1) do
      {:ok, [%{role: :assistant, content: content}]} when not is_nil(content) ->
        {:ok, state, content}

      other ->
        {:error, {:max_iterations_exceeded, other}}
    end
  end

  defp execute_streaming_loop(state, max_iterations) do
    case execute_streaming_step(state) do
      {:continue, updated_state} ->
        execute_streaming_loop(updated_state, max_iterations)

      {:done, content, updated_state} ->
        {:ok, updated_state, content}

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

  defp execute_streaming_step(state) do
    on_chunk = Keyword.get(state.opts, :on_chunk, fn _chunk -> :ok end)
    on_phase = Keyword.get(state.opts, :on_phase, fn _phase -> :ok end)
    on_tool_progress = Keyword.get(state.opts, :on_tool_progress, fn _info -> :ok end)

    with {:ok, messages} <- Memory.get_context(state.session_id, context_opts(state)),
         {:ok, stream} <- call_llm_stream(state, messages),
         {:ok, response} <-
           StreamConsumer.consume(stream, on_chunk: on_chunk, on_phase: on_phase) do
      updated_state = update_usage(state, response.usage)

      if is_binary(response.content) and response.content != "" do
        # Stream produced content -> final text response (streamed in real-time)
        Memory.add_message(state.session_id, :assistant, response.content)
        {:done, response.content, updated_state}
      else
        # Stream produced empty content -> LLM likely returned tool_calls
        # that the stream parser skipped. Fall back to blocking generate.
        Logger.debug("Stream empty, falling back to generate for tool calls")
        execute_step_with_tool_progress(state, on_tool_progress)
      end
    else
      {:error, reason} ->
        # stream_generate failed, fall back to blocking generate
        Logger.warning("stream_generate failed (#{inspect(reason)}), falling back to generate")

        execute_step_with_tool_progress(state, on_tool_progress)
    end
  end

  defp execute_step_with_tool_progress(state, on_tool_progress) do
    debug? = Application.get_env(:ragex, :debug_ai_responses, false)

    with {:ok, messages} <- Memory.get_context(state.session_id, context_opts(state)),
         {:ok, response} <- call_llm(state, messages) do
      updated_state = update_usage(state, response.usage)

      case response.tool_calls do
        nil ->
          content = response.content || ""
          if debug?, do: print_assistant_response(content, state.iterations)
          Memory.add_message(state.session_id, :assistant, content)
          {:done, content, updated_state}

        [] ->
          content = response.content || ""
          if debug?, do: print_assistant_response(content, state.iterations)
          Memory.add_message(state.session_id, :assistant, content)
          {:done, content, updated_state}

        tool_calls when is_list(tool_calls) ->
          Logger.debug("Agent making #{length(tool_calls)} tool call(s)")
          if debug?, do: print_tool_calls(tool_calls, response.content, state.iterations)

          Memory.add_message(state.session_id, :assistant, response.content || "",
            tool_calls: tool_calls
          )

          # Notify tool progress
          Enum.each(tool_calls, fn tc ->
            on_tool_progress.(%{
              type: :tool_call,
              name: tc.name,
              iteration: state.iterations
            })
          end)

          # Execute tools with dedup
          results = execute_tool_calls(tool_calls, state)

          # Check if ALL calls were repeats -> force text response
          all_repeated =
            Enum.all?(results, fn {_tc, result} ->
              match?({:ok, %{repeated: true}}, result)
            end)

          if all_repeated do
            Enum.each(results, fn {tool_call, result} ->
              result_str = format_tool_result(tool_call.name, result)
              if debug?, do: print_tool_result(tool_call.name, result_str)

              Memory.add_message(state.session_id, :tool, result_str,
                tool_call_id: tool_call.id,
                name: tool_call.name
              )
            end)

            force_text_response(updated_state)
          else
            Enum.each(results, fn {tool_call, result} ->
              result_str = format_tool_result(tool_call.name, result)
              if debug?, do: print_tool_result(tool_call.name, result_str)

              Memory.add_message(state.session_id, :tool, result_str,
                tool_call_id: tool_call.id,
                name: tool_call.name
              )

              Memory.add_tool_result(state.session_id, tool_call.id, result)
            end)

            updated_state =
              updated_state
              |> update_tool_history(tool_calls)
              |> Map.put(:iterations, updated_state.iterations + 1)
              |> Map.put(:tool_calls_made, updated_state.tool_calls_made + length(tool_calls))

            {:continue, updated_state}
          end
      end
    end
  end

  # --- Non-streaming execution loop ---

  defp execute_loop(state, max_iterations) when state.iterations >= max_iterations do
    Logger.warning("Agent reached max iterations (#{max_iterations})")
    # Return last assistant message or error
    case Memory.get_messages(state.session_id, limit: 1) do
      {:ok, [%{role: :assistant, content: content}]} when not is_nil(content) ->
        {:ok, state, content}

      other ->
        {:error, {:max_iterations_exceeded, other}}
    end
  end

  defp execute_loop(state, max_iterations) do
    case execute_step(state) do
      {:continue, updated_state} ->
        execute_loop(updated_state, max_iterations)

      {:done, content, updated_state} ->
        {:ok, updated_state, content}

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

  defp execute_step(state) do
    debug? = Application.get_env(:ragex, :debug_ai_responses, false)

    with {:ok, messages} <- Memory.get_context(state.session_id, context_opts(state)),
         {:ok, response} <- call_llm(state, messages) do
      # Update usage tracking
      updated_state = update_usage(state, response.usage)

      case response.tool_calls do
        nil ->
          # No tool calls - we're done
          content = response.content || ""
          if debug?, do: print_assistant_response(content, state.iterations)
          # Save assistant response
          Memory.add_message(state.session_id, :assistant, content)
          {:done, content, updated_state}

        [] ->
          # Empty tool calls - we're done
          content = response.content || ""
          if debug?, do: print_assistant_response(content, state.iterations)
          Memory.add_message(state.session_id, :assistant, content)
          {:done, content, updated_state}

        tool_calls when is_list(tool_calls) ->
          # Execute tool calls
          Logger.debug("Agent making #{length(tool_calls)} tool call(s)")
          if debug?, do: print_tool_calls(tool_calls, response.content, state.iterations)

          # Save assistant message with tool calls
          Memory.add_message(state.session_id, :assistant, response.content || "",
            tool_calls: tool_calls
          )

          # Execute each tool (with dedup) and add results
          results = execute_tool_calls(tool_calls, state)

          # Check if ALL calls were repeats -> force text response
          all_repeated =
            Enum.all?(results, fn {_tc, result} ->
              match?({:ok, %{repeated: true}}, result)
            end)

          if all_repeated do
            # Add the repeat messages to conversation
            Enum.each(results, fn {tool_call, result} ->
              result_str = format_tool_result(tool_call.name, result)
              if debug?, do: print_tool_result(tool_call.name, result_str)

              Memory.add_message(state.session_id, :tool, result_str,
                tool_call_id: tool_call.id,
                name: tool_call.name
              )
            end)

            # Force a text response without tools
            force_text_response(updated_state)
          else
            # Normal: add tool results and continue
            Enum.each(results, fn {tool_call, result} ->
              result_str = format_tool_result(tool_call.name, result)
              if debug?, do: print_tool_result(tool_call.name, result_str)

              Memory.add_message(state.session_id, :tool, result_str,
                tool_call_id: tool_call.id,
                name: tool_call.name
              )

              Memory.add_tool_result(state.session_id, tool_call.id, result)
            end)

            # Update state with history
            updated_state =
              updated_state
              |> update_tool_history(tool_calls)
              |> Map.put(:iterations, updated_state.iterations + 1)
              |> Map.put(:tool_calls_made, updated_state.tool_calls_made + length(tool_calls))

            {:continue, updated_state}
          end
      end
    end
  end

  defp call_llm(state, messages) do
    opts = [
      tools: state.tools,
      temperature: Keyword.get(state.opts, :temperature, @default_temperature),
      max_tokens: Keyword.get(state.opts, :max_tokens, @default_max_tokens),
      tool_choice: Keyword.get(state.opts, :tool_choice, "auto")
    ]

    # Format messages for provider
    formatted_messages = format_messages_for_provider(messages, state.provider_name)

    # Build prompt from messages
    {extracted_prompt, user_messages} =
      extract_system_prompt(formatted_messages, state.session_id)

    # Allow overriding the system prompt (e.g., for chat follow-up questions
    # where the report generation prompt should be replaced)
    system_prompt = Keyword.get(state.opts, :system_prompt_override, extracted_prompt)

    # Call provider
    prompt = build_prompt_from_messages(user_messages)

    state.provider.generate(prompt, %{messages: user_messages}, [
      {:system_prompt, system_prompt} | opts
    ])
  end

  defp call_llm_stream(state, messages) do
    opts = [
      tools: state.tools,
      temperature: Keyword.get(state.opts, :temperature, @default_temperature),
      max_tokens: Keyword.get(state.opts, :max_tokens, @default_max_tokens),
      tool_choice: Keyword.get(state.opts, :tool_choice, "auto")
    ]

    formatted_messages = format_messages_for_provider(messages, state.provider_name)

    {extracted_prompt, user_messages} =
      extract_system_prompt(formatted_messages, state.session_id)

    system_prompt = Keyword.get(state.opts, :system_prompt_override, extracted_prompt)
    prompt = build_prompt_from_messages(user_messages)

    state.provider.stream_generate(prompt, %{messages: user_messages}, [
      {:system_prompt, system_prompt} | opts
    ])
  end

  defp format_messages_for_provider(messages, :anthropic) do
    # Anthropic already handled by Memory.get_context
    messages
  end

  defp format_messages_for_provider(messages, _provider) do
    # OpenAI/DeepSeek format
    messages
  end

  defp extract_system_prompt(messages, session_id) do
    case Enum.split_with(messages, &(&1["role"] == "system" or &1[:role] == "system")) do
      {[system | _], others} ->
        {system["content"] || system[:content], others}

      {[], others} ->
        {default_system_prompt(get_project_path(session_id)), others}
    end
  end

  defp get_project_path(session_id) do
    case Memory.get_session(session_id) do
      {:ok, session} -> session.metadata[:project_path]
      _ -> nil
    end
  end

  defp build_prompt_from_messages(messages) do
    # Get the last user message as the prompt
    case Enum.reverse(messages) do
      [%{"content" => content} | _] when is_binary(content) -> content
      [%{content: content} | _] when is_binary(content) -> content
      _ -> ""
    end
  end

  defp default_system_prompt(project_path) do
    path_constraint =
      if project_path do
        """

        PROJECT CONTEXT:
        The project being analyzed is located at: #{project_path}
        CRITICAL: Any tool call that requires a "path" parameter MUST use exactly this path: #{project_path}
        Do NOT use ".", relative paths, parent directories, or any other path.
        """
      else
        ""
      end

    """
    You are an expert code analysis agent. You have access to tools for analyzing codebases.
    #{path_constraint}
    IMPORTANT RULES:
    1. Tool results are returned as JSON. Parse and use them directly.
    2. NEVER re-call a tool you already received results from. The data is already in the conversation.
    3. After receiving tool results, synthesize them into your final answer immediately.
    4. If a tool returns an error, explain the error and move on. Do NOT retry the same call.
    5. Use at most 3-5 tool calls total before producing your final answer.
    6. Your final response must be plain text (Markdown), NOT a tool call.

    When analyzing code:
    - Use semantic_search and hybrid_search for finding relevant code
    - Use find_dead_code, find_duplicates, scan_security for issues
    - Use analyze_quality, analyze_dependencies for structure
    - Use suggest_refactorings for improvements

    Always explain your findings clearly and provide specific recommendations.
    Work with the data you have. Do not loop trying to get more data.
    """
  end

  defp execute_tool_calls(tool_calls, state) do
    history = Map.get(state, :tool_call_history, %{})

    Enum.map(tool_calls, fn tool_call ->
      call_key = {tool_call.name, :erlang.phash2(tool_call.arguments)}

      if Map.has_key?(history, call_key) do
        Logger.warning("Skipping repeated tool call: #{tool_call.name}")

        {tool_call,
         {:ok,
          %{
            repeated: true,
            message:
              "This tool was already called with identical arguments. " <>
                "The result has NOT changed. Use the data from the previous call. " <>
                "Do NOT call this tool again. Produce your final answer now."
          }}}
      else
        Logger.debug("Executing tool: #{tool_call.name}")
        result = execute_single_tool(tool_call.name, tool_call.arguments)
        {tool_call, result}
      end
    end)
  end

  defp update_tool_history(state, tool_calls) do
    new_history =
      Enum.reduce(tool_calls, state.tool_call_history, fn tc, acc ->
        Map.put(acc, {tc.name, :erlang.phash2(tc.arguments)}, true)
      end)

    %{state | tool_call_history: new_history}
  end

  # When ALL tool calls in an iteration are repeats, force a text response
  # by calling the LLM without tool definitions.
  defp force_text_response(state) do
    Logger.warning("All tool calls repeated. Forcing text response without tools.")

    # Inject a directive so the AI synthesises already-available data rather
    # than reporting tool-call limitations as an audit failure.
    :ok =
      Memory.add_message(
        state.session_id,
        :user,
        "All necessary data is now available in the conversation from the analysis " <>
          "summary and previous tool call results. Write your complete final response " <>
          "now using all of that information. Do not mention tool call issues."
      )

    with {:ok, messages} <- Memory.get_context(state.session_id, context_opts(state)),
         {:ok, response} <- call_llm(%{state | tools: []}, messages) do
      content = response.content || ""
      Memory.add_message(state.session_id, :assistant, content)
      {:done, content, state}
    end
  end

  defp execute_single_tool(tool_name, arguments) do
    # Convert string keys to match MCP handler expectations
    args =
      arguments
      |> Enum.map(fn
        {k, v} when is_atom(k) -> {Atom.to_string(k), v}
        {k, v} -> {k, v}
      end)
      |> Enum.into(%{})

    try do
      case MCPTools.call_tool(tool_name, args) do
        {:ok, result} ->
          {:ok, result}

        {:error, reason} ->
          Logger.warning("Tool #{tool_name} failed: #{inspect(reason)}")
          {:error, reason}
      end
    rescue
      e ->
        Logger.error("Tool #{tool_name} raised exception: #{inspect(e)}")
        {:error, {:exception, Exception.message(e)}}
    catch
      kind, value ->
        Logger.error("Tool #{tool_name} threw #{kind}: #{inspect(value)}")
        {:error, {kind, value}}
    end
  end

  defp format_tool_result(tool_name, result) do
    json_body = encode_tool_result(result)

    """
    [TOOL_RESULT: #{tool_name}]
    The following is the complete JSON result from the "#{tool_name}" tool.
    Use this data directly in your response. Do NOT call this tool again.

    #{json_body}
    [/TOOL_RESULT]
    """
  end

  defp encode_tool_result({:ok, result}) when is_map(result) do
    result_str =
      with %{error_details: [_ | _]} = result <- result,
           do: update_in(result, [:error_details, Access.all(), :reason], &format_reason/1)

    case Jason.encode(result_str, pretty: true) do
      {:ok, encoded} -> encoded
      {:error, _} -> inspect(result_str, limit: :infinity, pretty: true)
    end
  end

  defp encode_tool_result({:ok, result}) do
    inspect(result, limit: :infinity, pretty: true)
  end

  defp encode_tool_result({:error, reason}) do
    Jason.encode!(%{error: true, reason: inspect(reason)})
  end

  # Console output helpers

  defp print_assistant_response(content, iteration) do
    IO.puts("")
    IO.puts(Colors.bold("--- Assistant (iteration #{iteration}) ---"))
    IO.puts(content)
    IO.puts(Colors.muted("--- end ---"))
    IO.puts("")
  end

  defp print_tool_calls(tool_calls, thinking, iteration) do
    IO.puts("")
    IO.puts(Colors.bold("--- Assistant (iteration #{iteration}) - Tool Calls ---"))

    if thinking && thinking != "" do
      IO.puts(Colors.muted("Thinking: #{thinking}"))
    end

    Enum.each(tool_calls, fn tc ->
      args_str =
        case Jason.encode(tc.arguments, pretty: true) do
          {:ok, s} -> s
          _ -> inspect(tc.arguments)
        end

      IO.puts(Colors.info("  -> #{tc.name}(#{args_str})"))
    end)

    IO.puts(Colors.muted("--- end ---"))
  end

  defp print_tool_result(tool_name, result_str) do
    IO.puts("")
    IO.puts(Colors.bold("--- Tool Result: #{tool_name} ---"))
    IO.puts(result_str)
    IO.puts(Colors.muted("--- end ---"))
    IO.puts("")
  end

  defp get_provider(opts) do
    case Keyword.get(opts, :provider) do
      nil -> Config.provider()
      :deepseek_r1 -> Ragex.AI.Provider.DeepSeekR1
      :openai -> Ragex.AI.Provider.OpenAI
      :anthropic -> Ragex.AI.Provider.Anthropic
      :ollama -> Ragex.AI.Provider.Ollama
      module when is_atom(module) -> module
    end
  end

  defp get_provider_name(opts) do
    case Keyword.get(opts, :provider) do
      nil -> Config.provider_name()
      name when is_atom(name) -> name
    end
  end

  defp get_tools(provider_name, opts) do
    case Keyword.get(opts, :tools) do
      nil -> ToolSchema.tools_for_provider(provider_name)
      tools when is_list(tools) -> tools
    end
  end

  defp update_usage(state, nil), do: state

  defp update_usage(state, usage) when is_map(usage) do
    current = state.total_usage

    updated = %{
      prompt_tokens:
        (current[:prompt_tokens] || 0) + (usage[:prompt_tokens] || usage["prompt_tokens"] || 0),
      completion_tokens:
        (current[:completion_tokens] || 0) +
          (usage[:completion_tokens] || usage["completion_tokens"] || 0),
      total_tokens:
        (current[:total_tokens] || 0) + (usage[:total_tokens] || usage["total_tokens"] || 0)
    }

    %{state | total_usage: updated}
  end
end