ex_claw

defmodule Mix.Tasks.ExClaw.Chat do
  @moduledoc """
  Starts an interactive chat REPL with an ADK Elixir agent.

      mix ex_claw.chat

  Options:
    --model       LLM model to use (default: gemini-2.0-flash)
    --name        Agent name (default: ex_claw)
    --skills-dir  Directory containing skills (default: ./skills/)
    --config      Path to a JSON config file (for tool policies, etc.)
  """
  @shortdoc "Interactive chat with an ExClaw agent"

  use Mix.Task

  @impl Mix.Task
  def run(args) do
    Mix.Task.run("app.start")

    {opts, _, _} =
      OptionParser.parse(args,
        strict: [model: :string, name: :string, skills_dir: :string, config: :string]
      )

    model = opts[:model] || "gemini-2.0-flash"
    name = opts[:name] || "ex_claw"
    skills_dir = opts[:skills_dir] || "./skills/"

    IO.puts("ExClaw Chat — #{name} (#{model})")

    config = load_config(opts[:config])
    policy = extract_tool_policy(config)
    context_guard_config = if config, do: ExClaw.Config.context_guard(config), else: %{}
    {:ok, skills} = ExClaw.Skill.Loader.scan(skills_dir)
    skills = ExClaw.ToolPolicy.filter(policy, skills)

    if skills != [] do
      skill_names = Enum.map_join(skills, ", ", & &1.name)
      IO.puts("Skills loaded: #{skill_names}")
    end

    IO.puts("Type 'exit' or Ctrl-C to quit.\n")

    base_instruction = """
    You are ExClaw, a helpful AI assistant built on ADK Elixir.
    Be concise and helpful. You're running as a CLI chat agent.
    """

    agent =
      ADK.Agent.LlmAgent.new(
        name: name,
        model: model,
        instruction: base_instruction
      )

    runner = ADK.Runner.new(app_name: name, agent: agent)

    loop(
      runner,
      name,
      "user-1",
      "chat-#{System.unique_integer([:positive])}",
      skills,
      base_instruction,
      model,
      context_guard_config
    )
  end

  defp loop(runner, name, user_id, session_id, skills, base_instruction, model, cg_config) do
    case IO.gets("you> ") do
      :eof ->
        IO.puts("\nBye!")

      input when is_binary(input) ->
        input = String.trim(input)

        if input in ["exit", "quit"] do
          IO.puts("Bye!")
        else
          cond do
            input == "" ->
              loop(runner, name, user_id, session_id, skills, base_instruction, model, cg_config)

            String.starts_with?(input, "/spawn ") ->
              task_text = String.trim_leading(input, "/spawn ")
              handle_spawn(task_text, name, model, base_instruction)
              loop(runner, name, user_id, session_id, skills, base_instruction, model, cg_config)

            input == "/sessions" ->
              handle_sessions()
              loop(runner, name, user_id, session_id, skills, base_instruction, model, cg_config)

            input == "/usage" ->
              handle_usage(session_id)
              loop(runner, name, user_id, session_id, skills, base_instruction, model, cg_config)

            true ->
              runner = maybe_apply_skill(runner, skills, input, name, base_instruction, model)

              case ADK.Runner.run(runner, user_id, session_id, input) do
                {:ok, events} ->
                  events
                  |> Enum.filter(&match?(%ADK.Event{author: a} when a != "user", &1))
                  |> Enum.each(fn event ->
                    text = extract_text(event)
                    if text && text != "", do: IO.puts("\n#{name}> #{text}\n")
                  end)

                  loop(
                    runner,
                    name,
                    user_id,
                    session_id,
                    skills,
                    base_instruction,
                    model,
                    cg_config
                  )

                {:error, reason} ->
                  IO.puts("\n[error] #{inspect(reason)}\n")

                  loop(
                    runner,
                    name,
                    user_id,
                    session_id,
                    skills,
                    base_instruction,
                    model,
                    cg_config
                  )
              end
          end
        end
    end
  end

  defp handle_spawn(task_text, name, model, instruction) do
    IO.puts("[spawning sub-agent...]")

    case ExClaw.SessionManager.spawn_session(task_text,
           agent_name: name,
           model: model,
           instruction: instruction
         ) do
      {:ok, session_id} ->
        IO.puts("[session #{session_id} started]")

        case ExClaw.SessionManager.await(session_id) do
          {:ok, result} ->
            IO.puts("\n#{name} (sub-agent)> #{result}\n")

          {:error, :timeout} ->
            IO.puts("\n[session #{session_id} timed out]\n")
        end
    end
  end

  defp handle_usage(session_id) do
    alias ExClaw.Usage.{Tracker, Budget, Forecaster, Cost}

    total = Tracker.total()

    IO.puts(
      "  Tokens — input: #{total.input_tokens}, output: #{total.output_tokens}, total: #{total.input_tokens + total.output_tokens}"
    )

    session_cost = Cost.session_cost(session_id)
    IO.puts("  Session cost: $#{:erlang.float_to_binary(session_cost, decimals: 6)}")

    budget_config = %{max_tokens: 1_000_000, warning_threshold: 0.8}

    case Budget.check(nil, budget_config) do
      :ok -> IO.puts("  Budget: OK")
      {:warning, pct} -> IO.puts("  Budget: WARNING (#{trunc(pct * 100)}% used)")
      {:exceeded, pct} -> IO.puts("  Budget: EXCEEDED (#{trunc(pct * 100)}% used)")
    end

    case Forecaster.burn_rate() do
      :no_data -> IO.puts("  Burn rate: insufficient data")
      {:ok, rate} -> IO.puts("  Burn rate: #{trunc(rate)} tokens/min")
    end
  rescue
    _ -> IO.puts("  [usage tracker not started]")
  end

  defp handle_sessions do
    sessions = ExClaw.SessionManager.list()

    if sessions == [] do
      IO.puts("[no sessions]")
    else
      Enum.each(sessions, fn s ->
        IO.puts(
          "  #{s.id} | #{s.status} | #{s.agent_name} | #{DateTime.to_iso8601(s.created_at)}"
        )
      end)
    end
  end

  defp maybe_apply_skill(runner, skills, input, name, base_instruction, model) do
    case ExClaw.Skill.match(skills, input) do
      nil ->
        runner

      %{instruction: nil} ->
        runner

      %{instruction: content} ->
        instruction = content <> "\n\n" <> base_instruction

        agent =
          ADK.Agent.LlmAgent.new(
            name: name,
            model: model,
            instruction: instruction
          )

        ADK.Runner.new(app_name: name, agent: agent)
    end
  end

  defp extract_text(%ADK.Event{content: %{parts: parts}}) when is_list(parts) do
    parts
    |> Enum.map(fn
      %{text: t} -> t
      _ -> nil
    end)
    |> Enum.reject(&is_nil/1)
    |> Enum.join("")
  end

  defp extract_text(_), do: nil

  defp load_config(nil), do: nil

  defp load_config(config_path) do
    case ExClaw.Config.load(config_path) do
      {:ok, config} ->
        config

      {:error, reason} ->
        IO.puts("[warn] Failed to load config: #{inspect(reason)}")
        nil
    end
  end

  defp extract_tool_policy(nil), do: nil

  defp extract_tool_policy(config) do
    case ExClaw.Config.tool_policy(config) do
      {:ok, policy} ->
        policy

      {:error, reason} ->
        IO.puts("[warn] Failed to load tool policy: #{inspect(reason)}")
        nil
    end
  end
end