Mix.Task.run("app.start")

Code.require_file(Path.expand("support/example_helper.exs", __DIR__))

alias CodexExamples.Support

Support.init!()

alias Codex.{Error, Events, Items, RunResultStreaming, Thread, TransportError}
alias Codex.ExamplesSupport

defmodule LiveWebSearchModes do
  @moduledoc false

  @default_prompt """
  Use web search to determine the latest stable Elixir release.
  Include at least one source URL in the answer.
  """
  @modes [:disabled, :cached, :live]
  @retry_prompt_suffix """
  You must use web search for this answer.
  """

  def main(args) do
    {modes, prompt} = parse_args(args)

    if ExamplesSupport.ollama_mode?() do
      IO.puts("""
      Ollama mode detected. This local OSS route does not guarantee Codex web-search events.
      Running a request-plumbing demo instead of enforcing live web-search event assertions.
      """)
    end

    failures = Enum.flat_map(modes, &run_mode(&1, prompt))

    if failures != [] do
      IO.puts("\nWeb search mode validation failures:")
      Enum.each(failures, fn failure -> IO.puts("  - #{inspect(failure)}") end)
      System.halt(1)
    end
  end

  defp run_mode(mode, prompt) do
    IO.puts("\n--- web_search_mode=#{mode} ---")

    if mode == :disabled do
      IO.puts("Disabled mode is passed explicitly (web_search=\"disabled\").")
    end

    case run_mode_once(mode, prompt) do
      {:ok, final_state} ->
        report_final_state(mode, final_state)
        []

      {:retry_required, reason} ->
        retry_prompt = String.trim("#{prompt}\n\n#{@retry_prompt_suffix}")
        IO.puts("Retrying with stricter prompt due to missing web search events.")
        IO.puts("Retry reason: #{inspect(reason)}")

        case run_mode_once(mode, retry_prompt) do
          {:ok, final_state} ->
            report_final_state(mode, final_state)
            []

          {:retry_required, retry_reason} ->
            [{mode, retry_reason}]

          {:error, retry_error} ->
            [{mode, retry_error}]
        end

      {:error, reason} ->
        [{mode, reason}]
    end
  end

  defp run_mode_once(mode, prompt) do
    with {:ok, codex_opts} <- Support.codex_options(%{}),
         {:ok, thread_opts} <- Codex.Thread.Options.new(%{web_search_mode: mode}),
         {:ok, thread} <- Codex.start_thread(codex_opts, Support.thread_opts!(thread_opts)),
         {:ok, result} <- Thread.run_streamed(thread, prompt),
         {:ok, final_state} <- consume_result(result) do
      validate_mode_expectation(mode, final_state)
    else
      {:error, reason} ->
        IO.puts("Failed to start streamed turn: #{inspect(reason)}")
        {:error, reason}
    end
  end

  defp consume_result(result) do
    try do
      final_state =
        result
        |> RunResultStreaming.raw_events()
        |> Enum.reduce(%{web_search?: false, final_response: nil}, &handle_event/2)

      {:ok, final_state}
    rescue
      error in [Error, TransportError] ->
        render_transport_error(error)
        {:error, error}
    end
  end

  defp validate_mode_expectation(:disabled, %{web_search?: true}) do
    {:error, {:unexpected_web_search_events, :observed_in_disabled_mode}}
  end

  defp validate_mode_expectation(:disabled, final_state), do: {:ok, final_state}

  defp validate_mode_expectation(:live, final_state) do
    if ExamplesSupport.ollama_mode?() do
      {:ok, final_state}
    else
      validate_live_mode_expectation(final_state)
    end
  end

  # Cached mode only permits cached search results. A turn may legitimately emit
  # no web-search events when no cached result is available or the model answers
  # without using the tool.
  defp validate_mode_expectation(:cached, final_state), do: {:ok, final_state}

  defp validate_live_mode_expectation(%{web_search?: true} = final_state),
    do: {:ok, final_state}

  defp validate_live_mode_expectation(_final_state),
    do: {:retry_required, {:expected_web_search_events, :none_observed}}

  defp report_final_state(mode, final_state) do
    if final_state.web_search? do
      IO.puts("Observed web search events.")
    else
      report_no_web_search_events(mode)
    end

    IO.puts("Final response (illustrative only): #{final_state.final_response || "<none>"}")
  end

  defp report_no_web_search_events(:cached) do
    IO.puts(
      "No web search events observed. Cached mode may legitimately complete without a cache hit."
    )
  end

  defp report_no_web_search_events(:live) do
    if ExamplesSupport.ollama_mode?() do
      IO.puts(
        "No web search events observed. Ollama mode validates request plumbing only for live web_search."
      )
    else
      IO.puts("No web search events observed.")
    end
  end

  defp report_no_web_search_events(_mode), do: IO.puts("No web search events observed.")

  defp handle_event(%Events.ItemStarted{item: %Items.WebSearch{query: query}}, state) do
    IO.puts("Web search started: #{format_query(query)}")
    %{state | web_search?: true}
  end

  defp handle_event(%Events.ItemCompleted{item: %Items.WebSearch{query: query}}, state) do
    IO.puts("Web search completed: #{format_query(query)}")
    %{state | web_search?: true}
  end

  defp handle_event(%Events.ItemCompleted{item: %Items.AgentMessage{text: text}}, state) do
    %{state | final_response: text}
  end

  defp handle_event(%Events.TurnCompleted{final_response: response}, state) do
    %{state | final_response: extract_text(response) || state.final_response}
  end

  defp handle_event(_event, state), do: state

  defp extract_text(%Items.AgentMessage{text: text}) when is_binary(text), do: text
  defp extract_text(%{"type" => "text", "text" => text}) when is_binary(text), do: text
  defp extract_text(%{type: "text", text: text}) when is_binary(text), do: text
  defp extract_text(other) when is_binary(other), do: other
  defp extract_text(_), do: nil

  defp format_query(query) when is_binary(query) do
    case String.trim(query) do
      "" -> "<pending>"
      trimmed -> trimmed
    end
  end

  defp format_query(_query), do: "<pending>"

  defp parse_args([mode | rest]) when mode in ["disabled", "cached", "live"] do
    {[web_search_mode!(mode)], parse_prompt(rest)}
  end

  defp parse_args(args), do: {@modes, parse_prompt(args)}

  defp web_search_mode!("disabled"), do: :disabled
  defp web_search_mode!("cached"), do: :cached
  defp web_search_mode!("live"), do: :live

  defp parse_prompt([prompt | _]), do: prompt
  defp parse_prompt(_args), do: @default_prompt

  defp render_transport_error(%Error{} = error) do
    status = error.details[:exit_status]
    stderr = error.details[:stderr]

    IO.puts("""
    Failed to run codex#{format_exit_status(status)}.
    #{error.message}
    Ensure the codex CLI is installed on PATH and you're logged in (or set CODEX_API_KEY).
    stderr: #{String.trim(to_string(stderr || ""))}
    """)
  end

  defp render_transport_error(%TransportError{exit_status: status, stderr: stderr}) do
    IO.puts("""
    Failed to run codex (exit #{inspect(status)}).
    Ensure the codex CLI is installed on PATH and you're logged in (or set CODEX_API_KEY).
    stderr: #{String.trim(to_string(stderr || ""))}
    """)
  end

  defp format_exit_status(nil), do: ""
  defp format_exit_status(status), do: " (exit #{inspect(status)})"
end

LiveWebSearchModes.main(System.argv())