"""
Gemini Agent 基础示例

使用 Gemini 2.5 Pro 模型，演示带工具调用的 Agent

依赖:
    pip install httpx python-dotenv

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

import os
import sys
import json
import asyncio
from typing import Dict, Any, List, Optional
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 tool, ToolResult, get_tool_registry
from agent.runner import AgentRunner
from agent.llm.providers.gemini import create_gemini_llm_call


# ============================================================
# 定义工具
# ============================================================

@tool()
async def get_current_weather(location: str, unit: str = "celsius", uid: str = "") -> Dict[str, Any]:
    """
    获取指定地点的当前天气

    Args:
        location: 城市名称，如 "北京"、"San Francisco"
        unit: 温度单位，"celsius" 或 "fahrenheit"

    Returns:
        天气信息字典
    """
    # 模拟天气数据
    weather_data = {
        "北京": {"temp": 15, "condition": "晴朗", "humidity": 45},
        "上海": {"temp": 20, "condition": "多云", "humidity": 60},
        "San Francisco": {"temp": 18, "condition": "Foggy", "humidity": 70},
        "New York": {"temp": 10, "condition": "Rainy", "humidity": 80}
    }

    data = weather_data.get(location, {"temp": 22, "condition": "Unknown", "humidity": 50})

    if unit == "fahrenheit":
        data["temp"] = data["temp"] * 9/5 + 32

    return {
        "location": location,
        "temperature": data["temp"],
        "unit": unit,
        "condition": data["condition"],
        "humidity": data["humidity"]
    }


@tool()
async def calculate(expression: str, uid: str = "") -> ToolResult:
    """
    执行数学计算

    Args:
        expression: 数学表达式，如 "2 + 2"、"10 * 5"

    Returns:
        计算结果
    """
    try:
        # 安全地计算简单表达式
        # 注意：实际生产环境应使用更安全的方法
        result = eval(expression, {"__builtins__": {}}, {})

        return ToolResult(
            title="计算结果",
            output=f"{expression} = {result}",
            long_term_memory=f"计算了 {expression}"
        )
    except Exception as e:
        return ToolResult(
            title="计算错误",
            output=f"无法计算 '{expression}': {str(e)}",
            long_term_memory=f"计算失败: {expression}"
        )


@tool()
async def search_knowledge(query: str, max_results: int = 3, uid: str = "") -> ToolResult:
    """
    搜索知识库

    Args:
        query: 搜索关键词
        max_results: 返回结果数量

    Returns:
        搜索结果
    """
    # 模拟知识库搜索
    knowledge_base = {
        "Python": "Python 是一种高级编程语言，以简洁易读的语法著称。",
        "Agent": "Agent 是能够感知环境并采取行动以实现目标的智能体。",
        "Gemini": "Gemini 是 Google 开发的多模态大语言模型系列。",
        "AI": "人工智能(AI)是计算机科学的一个分支，致力于创建智能机器。"
    }

    results = []
    for key, value in knowledge_base.items():
        if query.lower() in key.lower() or query.lower() in value.lower():
            results.append({"title": key, "content": value})
            if len(results) >= max_results:
                break

    if not results:
        output = f"未找到关于 '{query}' 的信息"
    else:
        output = "\n\n".join([f"**{r['title']}**\n{r['content']}" for r in results])

    return ToolResult(
        title=f"搜索结果: {query}",
        output=output,
        long_term_memory=f"搜索了 '{query}'，找到 {len(results)} 条结果"
    )


# ============================================================
# 主函数
# ============================================================

async def main():
    print("=" * 60)
    print("Gemini Agent 基础示例")
    print("=" * 60)
    print()

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

    # 打印可用工具
    print("可用工具:")
    for tool_name in registry.get_tool_names():
        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 = "北京今天的天气怎么样？顺便帮我计算一下 15 * 8 等于多少。"

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

    # 运行 Agent
    async for event in runner.run(
        task=task,
        model="gemini-2.5-pro",
        tools=["get_current_weather", "calculate", "search_knowledge"],
        max_iterations=5,
        enable_memory=False,  # 暂不启用记忆
        auto_execute_tools=True,
        system_prompt="你是一个有用的AI助手，可以使用工具来帮助用户。请简洁明了地回答问题。"
    ):
        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":
            print(f"   结果: {data['result'][:100]}...")
            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())