"""
Subagent 工具真实测试

使用真实 LLM 测试 subagent 工具的三种模式：
1. delegate - 委托子任务
2. explore - 并行探索方案
3. evaluate - 评估结果
"""

import os
import sys
import asyncio
from pathlib import Path

# 添加项目根目录到 Python 路径
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from dotenv import load_dotenv
load_dotenv()

from agent.llm.prompts import SimplePrompt
from agent.core.runner import AgentRunner
from agent.trace import (
    FileSystemTraceStore,
    Trace,
    Message,
)
from agent.llm import create_openrouter_llm_call


async def main():
    # 路径配置
    base_dir = Path(__file__).parent
    project_root = base_dir.parent.parent
    prompt_path = base_dir / "test.prompt"
    output_dir = base_dir / "output"
    output_dir.mkdir(exist_ok=True)

    print("=" * 60)
    print("Subagent 工具测试 (真实 LLM)")
    print("=" * 60)
    print()

    # 1. 加载 prompt
    print("1. 加载 prompt...")
    prompt = SimplePrompt(prompt_path)

    # 提取配置
    system_prompt = prompt._messages.get("system", "")
    user_task = prompt._messages.get("user", "")
    model_name = prompt.config.get('model', 'gemini-2.5-flash')
    temperature = float(prompt.config.get('temperature', 0.3))

    print(f"   - 任务: {user_task[:80]}...")
    print(f"   - 模型: {model_name}")

    # 2. 构建消息
    print("2. 构建任务消息...")
    messages = prompt.build_messages()

    # 3. 创建 Agent Runner
    print("3. 创建 Agent Runner...")
    print(f"   - 模型: {model_name} (via OpenRouter)")

    # Trace 输出到测试目录
    trace_dir = base_dir / ".trace"
    trace_dir.mkdir(exist_ok=True)
    print(f"   - Trace 目录: {trace_dir}")

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(base_path=str(trace_dir)),
        llm_call=create_openrouter_llm_call(model=f"google/{model_name}"),
        skills_dir=None,
        debug=True
    )

    # 4. Agent 模式执行
    print(f"4. 启动 Agent 模式...")
    print()

    final_response = ""
    current_trace_id = None
    subagent_calls = []

    async for item in runner.run(
        task=user_task,
        messages=messages,
        system_prompt=system_prompt,
        model=f"google/{model_name}",
        temperature=temperature,
        max_iterations=30,  # 增加迭代次数以支持多个 subagent 调用
    ):
        # 处理 Trace 对象
        if isinstance(item, Trace):
            current_trace_id = item.trace_id
            if item.status == "running":
                print(f"[Trace] 开始: {item.trace_id[:8]}")
            elif item.status == "completed":
                print(f"[Trace] 完成")
                print(f"  - Total messages: {item.total_messages}")
                print(f"  - Total tokens: {item.total_tokens}")
                print(f"  - Total cost: ${item.total_cost:.4f}")
            elif item.status == "failed":
                print(f"[Trace] 失败: {item.error_message}")

        # 处理 Message 对象
        elif isinstance(item, Message):
            if item.role == "assistant":
                content = item.content
                if isinstance(content, dict):
                    text = content.get("text", "")
                    tool_calls = content.get("tool_calls")

                    if text and not tool_calls:
                        final_response = text
                        print(f"[Response] Agent 完成")
                    elif text:
                        print(f"[Assistant] {text[:100]}...")

                    if tool_calls:
                        for tc in tool_calls:
                            tool_name = tc.get("function", {}).get("name", "unknown")
                            print(f"[Tool Call] {tool_name}")

                            # 记录 subagent 调用
                            if tool_name == "subagent":
                                import json
                                args = tc.get("function", {}).get("arguments", {})
                                # arguments 可能是字符串，需要解析
                                if isinstance(args, str):
                                    try:
                                        args = json.loads(args)
                                    except:
                                        args = {}
                                mode = args.get("mode", "unknown")
                                subagent_calls.append({
                                    "mode": mode,
                                    "task": args.get("task", args.get("background", ""))[:50]
                                })
                                print(f"  → mode: {mode}")

            elif item.role == "tool":
                content = item.content
                if isinstance(content, dict):
                    tool_name = content.get("tool_name", "unknown")
                    print(f"[Tool Result] {tool_name}")
                if item.description:
                    desc = item.description[:80] if len(item.description) > 80 else item.description
                    print(f"  {desc}...")

    # 5. 输出结果
    print()
    print("=" * 60)
    print("Agent 响应:")
    print("=" * 60)
    print(final_response)
    print("=" * 60)
    print()

    # 6. 统计 subagent 调用
    print("=" * 60)
    print("Subagent 调用统计:")
    print("=" * 60)
    delegate_count = sum(1 for call in subagent_calls if call["mode"] == "delegate")
    explore_count = sum(1 for call in subagent_calls if call["mode"] == "explore")
    evaluate_count = sum(1 for call in subagent_calls if call["mode"] == "evaluate")

    print(f"  - delegate 模式: {delegate_count} 次")
    print(f"  - explore 模式: {explore_count} 次")
    print(f"  - evaluate 模式: {evaluate_count} 次")
    print(f"  - 总计: {len(subagent_calls)} 次")
    print()

    for i, call in enumerate(subagent_calls, 1):
        print(f"  {i}. [{call['mode']}] {call['task']}...")
    print("=" * 60)
    print()

    # 7. 保存结果
    output_file = output_dir / "subagent_test_result.txt"
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("=" * 60 + "\n")
        f.write("Agent 响应\n")
        f.write("=" * 60 + "\n\n")
        f.write(final_response)
        f.write("\n\n" + "=" * 60 + "\n")
        f.write("Subagent 调用统计\n")
        f.write("=" * 60 + "\n\n")
        f.write(f"delegate 模式: {delegate_count} 次\n")
        f.write(f"explore 模式: {explore_count} 次\n")
        f.write(f"evaluate 模式: {evaluate_count} 次\n")
        f.write(f"总计: {len(subagent_calls)} 次\n\n")
        for i, call in enumerate(subagent_calls, 1):
            f.write(f"{i}. [{call['mode']}] {call['task']}...\n")

    print(f"✓ 结果已保存到: {output_file}")
    print()

    # 8. 可视化提示
    print("=" * 60)
    print("Trace 信息:")
    print("=" * 60)
    print(f"Trace ID: {current_trace_id}")
    print(f"Trace 目录: {trace_dir}")
    print()
    print("查看 trace 文件:")
    print(f"   ls -la {trace_dir}")
    print()
    print("或启动 API Server 可视化:")
    print("   python3 api_server.py")
    print("   访问: http://localhost:8000/api/traces")
    print("=" * 60)


if __name__ == "__main__":
    asyncio.run(main())