"""
特征提取示例

使用 Agent 模式 + Skills + 多模态支持
"""

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, RunConfig
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"
    feature_md_path = base_dir / "input_1" / "feature.md"
    image_path = base_dir / "input_1" / "image.png"
    output_dir = base_dir / "output_1"
    output_dir.mkdir(exist_ok=True)

    # Skills 目录（可选：用户自定义 skills）
    # 注意：内置 skills（agent/skills/core.md）会自动加载
    skills_dir = None  # 或者指定自定义 skills 目录，如: project_root / "skills"

    print("=" * 60)
    print("特征提取任务 (Agent 模式)")
    print("=" * 60)
    print()

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

    # 提取 system prompt 和 user template
    system_prompt = prompt._messages.get("system", "")
    user_template = prompt._messages.get("user", "")

    # 2. 读取特征描述
    print("2. 读取特征描述...")
    with open(feature_md_path, 'r', encoding='utf-8') as f:
        feature_text = f.read()

    # 3. 构建任务文本（包含图片）
    print("3. 构建任务（文本 + 图片）...")

    # 使用 prompt 构建多模态消息
    temp_messages = prompt.build_messages(
        text=feature_text,
        images=image_path
    )

    # 提取用户消息（包含文本和图片）
    user_message_with_image = None
    for msg in temp_messages:
        if msg["role"] == "user":
            user_message_with_image = msg
            break

    if not user_message_with_image:
        raise ValueError("No user message found in prompt")

    print(f"   - 任务已构建（包含图片: {image_path.name}）")

    # 4. 创建 Agent Runner（配置 skills）
    print("4. 创建 Agent Runner...")
    print(f"   - Skills 目录: {skills_dir}")
    print(f"   - 模型: Claude Sonnet 4.5 (via OpenRouter)")

    runner = AgentRunner(
        trace_store=FileSystemTraceStore(base_path=".trace"),
        llm_call=create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5"),
        skills_dir=skills_dir,  # 可选：加载额外的用户自定义 skills（内置 skills 会自动加载）
        debug=True  # 启用 debug，输出到 .trace/
    )

    # 5. Agent 模式执行
    # 注意：使用 OpenRouter 时，模型在创建 llm_call 时已指定
    # 这里传入的 model 参数会被忽略（由 llm_call 内部控制）
    print(f"5. 启动 Agent 模式...")
    print()

    final_response = ""
    current_trace_id = None  # 保存 trace_id 用于后续测试

    async for item in runner.run(
        messages=[user_message_with_image],
        config=RunConfig(
            system_prompt=system_prompt,
            model="anthropic/claude-sonnet-4.5",
            temperature=float(prompt.config.get('temperature', 0.3)),
            max_iterations=1000,
            name="特征提取任务",
        ),
    ):
        # 处理 Trace 对象（整体状态变化）
        if isinstance(item, Trace):
            current_trace_id = item.trace_id  # 保存 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] 失败")

        # 处理 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}")

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

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

    # 7. 保存结果
    output_file = output_dir / "result.txt"
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(final_response)

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

    # 提示使用 API 可视化
    print("=" * 60)
    print("可视化 Step Tree:")
    print("=" * 60)
    print("1. 启动 API Server:")
    print("   python3 api_server.py")
    print()
    print("2. 浏览器访问:")
    print("   http://43.106.118.91:8000/api/traces")
    print()
    print(f"3. Trace ID: {current_trace_id}")
    print("=" * 60)


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