"""
百度搜索 Agent 示例

使用 browser-use 工具在百度搜索并返回结果

依赖:
    pip install playwright httpx python-dotenv
    playwright install chromium

使用方法:
    python examples/baidu_search_agent.py
"""

import os
import sys
import json
import asyncio
from dotenv import load_dotenv

# 添加项目根目录到 Python 路径
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 加载环境变量
load_dotenv()

# 导入框架
from agent.tools import get_tool_registry
from agent.runner import AgentRunner
from agent.llm.providers.gemini import create_gemini_llm_call

# 导入 browser-use 工具（这会自动注册工具）
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "tools"))
import browserUseTools


async def main():
    print("=" * 60)
    print("百度搜索 Agent 示例")
    print("=" * 60)
    print()

    # 获取工具注册表
    registry = get_tool_registry()

    # 打印可用工具
    print("可用的浏览器工具:")
    for tool_name in registry.get_tool_names():
        if "browser" in tool_name or "baidu" in tool_name:
            print(f"  - {tool_name}")
    print()

    # 创建 Gemini LLM 调用函数
    gemini_llm_call = create_gemini_llm_call()

    # 创建 Agent Runner
    runner = AgentRunner(
        tool_registry=registry,
        llm_call=gemini_llm_call,
    )

    # 测试任务：多步骤百度搜索
    task = """请完成以下任务：
1. 在百度搜索"瑜伽裤美女"
2. 返回第2页的搜索结果数据（标题、链接、摘要）

请使用 baidu_search 工具完成。"""

    print(f"任务: {task}")
    print("-" * 60)
    print()

    # 运行 Agent
    async for event in runner.run(
        task=task,
        model="gemini-2.5-pro",
        tools=[
            "baidu_search"
        ],
        max_iterations=5,
        enable_memory=False,
        auto_execute_tools=True,
        system_prompt="""你是一个有用的AI助手，可以使用浏览器工具来帮助用户完成网页操作任务。

可用工具说明：
- baidu_search: 在百度搜索并返回结果，支持 page 参数用于指定页码

请按照用户的要求，逐步使用这些工具完成任务。"""
    ):
        event_type = event.type
        data = event.data

        if event_type == "trace_started":
            print(f"✓ Trace 开始: {data['trace_id']}")
            print()

        elif event_type == "llm_call_completed":
            print(f"🤖 LLM 响应:")
            if data.get("content"):
                print(f"   {data['content']}")
            if data.get("tool_calls"):
                print(f"   工具调用: {len(data['tool_calls'])} 个")
            print(f"   Tokens: {data.get('tokens', 0)}")
            print()

        elif event_type == "tool_executing":
            print(f"🔧 执行工具: {data['tool_name']}")
            print(f"   参数: {json.dumps(data['arguments'], ensure_ascii=False)}")

        elif event_type == "tool_result":
            result_preview = data['result'][:200] if len(data['result']) > 200 else data['result']
            print(f"   结果预览: {result_preview}...")
            print()

        elif event_type == "conclusion":
            print(f"✅ 最终回答:")
            print(f"   {data['content']}")
            print()

        elif event_type == "trace_completed":
            print(f"✓ Trace 完成")
            print(f"   总 Tokens: {data.get('total_tokens', 0)}")
            print(f"   总成本: ${data.get('total_cost', 0):.4f}")
            print()

        elif event_type == "trace_failed":
            print(f"❌ Trace 失败: {data.get('error')}")
            print()


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