defmodule Mix.Tasks.ExAi.Agent.Reply do
  @shortdoc "Reply to a previous ExAI agent run"
  @moduledoc """
  Continues a prior agent conversation by loading persisted history from a run id.

      mix ex_ai.agent.reply <run_id> --input "follow up"
      mix ex_ai.agent.reply <run_id> --input "follow up" --json
  """

  use Mix.Task

  alias ExAI.CLI.AgentCatalog
  alias ExAI.CLI.CommandTelemetry
  alias ExAI.CLI.JSON
  alias ExAI.CLI.Output
  alias ExAI.CLI.Runtime
  alias ExAI.Error
  alias ExAI.Persistence.History
  alias ExAI.Persistence.Store

  @switches [input: :string, dir: :string, json: :boolean]
  @history_roles [:system, :user, :assistant, :tool]

  @impl true
  def run(argv) do
    {opts, args, invalid} = OptionParser.parse(argv, strict: @switches)
    opts = Map.new(opts)
    persistence_opts = Runtime.persistence_opts(opts)
    started_ms = CommandTelemetry.start("ex_ai.agent.reply")

    with :ok <- validate_invalid_options(invalid),
         {:ok, parent_run_id} <- parse_run_id(args),
         {:ok, input} <- parse_input(opts),
         :ok <- Runtime.boot(opts),
         {:ok, parent_run} <- Store.load_run(parent_run_id, persistence_opts),
         {:ok, agent_name} <- extract_agent_name(parent_run),
         {:ok, agent} <- AgentCatalog.fetch(agent_name),
         {:ok, history_messages} <- load_history_messages(parent_run_id, persistence_opts),
         {:ok, context, events} <-
           execute(agent, agent_name, parent_run_id, input, history_messages),
         {:ok, persisted} <- Runtime.persist_run_bundle(context, events, opts) do
      payload =
        JSON.success(%{
          run_id: context.run_id,
          parent_run_id: parent_run_id,
          status: context.status,
          result: context.result,
          usage: usage_from_result(context.result),
          metadata: context.metadata,
          persisted: %{
            events: persisted.events_persisted,
            messages: persisted.messages_persisted
          }
        })

      CommandTelemetry.stop("ex_ai.agent.reply", started_ms, %{run_id: context.run_id})
      emit_success(payload, opts)
    else
      {:error, %Error{} = error} ->
        CommandTelemetry.error("ex_ai.agent.reply", started_ms, error)
        emit_error(error, opts, %{})

      {:error, reason} ->
        error = Error.new(:validation_error, reason)
        CommandTelemetry.error("ex_ai.agent.reply", started_ms, error)
        emit_error(error, opts, %{})
    end
  end

  @spec execute(ExAI.Agent.t(), String.t(), String.t(), String.t(), [map()]) ::
          {:ok, ExAI.Graph.Context.t(), [ExAI.Graph.Event.t()]}
          | {:error, ExAI.Error.t()}
  defp execute(%ExAI.Agent{} = agent, agent_name, parent_run_id, input, history_messages) do
    metadata = %{
      agent_name: agent_name,
      parent_run_id: parent_run_id,
      cli_task: "ex_ai.agent.reply"
    }

    reply_agent =
      %ExAI.Agent{
        agent
        | history_processor: fn messages -> history_messages ++ messages end
      }

    runtime_opts = [metadata: metadata]

    case Runtime.run_agent(reply_agent, input, runtime_opts) do
      {:ok, context, events} ->
        {:ok, context, events}

      {:error, %Error{} = error, _context, _events} ->
        {:error, error}
    end
  end

  @spec extract_agent_name(map()) :: {:ok, String.t()} | {:error, ExAI.Error.t()}
  defp extract_agent_name(run) do
    metadata = Map.get(run, :metadata, %{})

    case Map.get(metadata, :agent_name) || Map.get(metadata, "agent_name") do
      nil ->
        {:error,
         Error.new(:validation_error, "run metadata does not contain agent_name",
           details: %{run_id: Map.get(run, :run_id)}
         )}

      agent_name ->
        {:ok, to_string(agent_name)}
    end
  end

  @spec load_history_messages(String.t(), keyword()) ::
          {:ok, [map()]} | {:error, ExAI.Error.t()}
  defp load_history_messages(run_id, opts) do
    with {:ok, records} <- History.list_messages(run_id, opts) do
      messages =
        records
        |> Enum.filter(fn record -> normalize_role(record.role) in @history_roles end)
        |> Enum.map(fn record ->
          %{
            role: normalize_role(record.role),
            content: record.content,
            metadata: record.metadata
          }
        end)

      {:ok, messages}
    end
  end

  @spec parse_run_id([String.t()]) :: {:ok, String.t()} | {:error, String.t()}
  defp parse_run_id([run_id | _]) when is_binary(run_id), do: {:ok, run_id}
  defp parse_run_id(_), do: {:error, "usage: mix ex_ai.agent.reply <run_id> --input <text>"}

  @spec parse_input(map()) :: {:ok, String.t()} | {:error, String.t()}
  defp parse_input(%{input: input}) when is_binary(input) and input != "", do: {:ok, input}
  defp parse_input(_), do: {:error, "--input is required"}

  @spec normalize_role(term()) :: atom() | String.t()
  defp normalize_role(role) when is_atom(role), do: role

  defp normalize_role(role) when is_binary(role) do
    case String.downcase(role) do
      "system" -> :system
      "user" -> :user
      "assistant" -> :assistant
      "tool" -> :tool
      other -> other
    end
  end

  defp normalize_role(role), do: role

  @spec usage_from_result(term()) :: term()
  defp usage_from_result(%{} = result), do: Map.get(result, :usage)
  defp usage_from_result(_), do: nil

  @spec validate_invalid_options([{String.t(), String.t() | nil}]) :: :ok | {:error, String.t()}
  defp validate_invalid_options([]), do: :ok

  defp validate_invalid_options(invalid) do
    {:error, "invalid options: #{inspect(invalid)}"}
  end

  @spec emit_success(map(), map()) :: :ok
  defp emit_success(payload, %{json: true}) do
    Output.emit_json(payload)
  end

  defp emit_success(payload, _opts) do
    Output.emit_info("run #{payload.run_id} completed with status #{payload.status}")
    :ok
  end

  @spec emit_error(ExAI.Error.t(), map(), map()) :: no_return()
  defp emit_error(%Error{} = error, opts, payload_overrides) do
    Output.halt_error(error, opts, payload_overrides)
  end
end