"""
Google ADK 自定义工具开发完整示例
展示各种工具开发技巧和最佳实践

对应教程：第04章 - 自定义工具开发
"""

# 导入 ADK 核心模块
from google.adk.agents import Agent          # Agent 类
from google.adk.tools import AgentTool       # AgentTool（Agent 作为工具）

# 导入类型提示
from typing import Optional, List            # 可选类型和列表类型

# 导入运行相关模块
from google.adk.runners import Runner        # 运行器
from google.adk.sessions import InMemorySessionService  # 会话服务
from google.genai import types               # 类型定义

# 导入异步库
import asyncio                               # 异步编程库


# ========================================
# 示例一：基础函数工具
# ========================================

def get_weather(city: str) -> dict:
    """
    获取指定城市的天气信息

    Args:
        city (str): 城市名称，例如"北京"、"上海"

    Returns:
        dict: 包含天气状态的字典
            - status: "success" 或 "error"
            - city: 城市名称
            - temperature: 温度
            - condition: 天气状况
    """
    # 模拟天气数据
    weather_data = {                        # 天气数据库
        "北京": {"temp": "25°C", "condition": "晴天"},  # 北京
        "上海": {"temp": "28°C", "condition": "多云"},  # 上海
        "广州": {"temp": "32°C", "condition": "雷阵雨"}, # 广州
    }

    # 查找城市天气
    data = weather_data.get(city)           # 获取天气数据

    if not data:                            # 如果找不到
        return {                            # 返回错误
            "status": "error",
            "error_message": f"未找到'{city}'的天气信息"
        }

    return {                                # 返回成功结果
        "status": "success",
        "city": city,
        "temperature": data["temp"],
        "condition": data["condition"],
    }


# ========================================
# 示例二：带可选参数的工具
# ========================================

def search_restaurant(
    city: str,                              # 必选参数
    cuisine: str = "中餐",                   # 可选参数，默认中餐
    price_range: str = "中等",               # 可选参数，默认中等
) -> dict:
    """
    搜索指定城市的餐厅

    Args:
        city (str): 城市名称
        cuisine (str, optional): 菜系类型，默认"中餐"
        price_range (str, optional): 价格范围，默认"中等"

    Returns:
        dict: 搜索结果
    """
    return {                                # 返回搜索结果
        "status": "success",
        "city": city,
        "cuisine": cuisine,
        "price_range": price_range,
        "results": [                        # 餐厅列表
            {"name": "美味餐厅", "rating": 4.5, "price": "人均100元"},
            {"name": "佳肴轩", "rating": 4.2, "price": "人均80元"},
        ],
    }


# ========================================
# 示例三：使用 Optional 类型
# ========================================

def create_user_profile(
    username: str,                          # 必选参数
    bio: Optional[str] = None,              # 可选参数
    avatar_url: Optional[str] = None,       # 可选参数
) -> dict:
    """
    创建用户档案

    Args:
        username (str): 用户名
        bio (str, optional): 个人简介
        avatar_url (str, optional): 头像 URL

    Returns:
        dict: 创建结果
    """
    profile = {                             # 创建档案
        "username": username,               # 用户名
        "bio": bio or "这个人很懒，什么都没写。",  # 简介
        "avatar_url": avatar_url or "",     # 头像
    }

    return {                                # 返回结果
        "status": "success",
        "profile": profile,
    }


# ========================================
# 示例四：List 类型参数
# ========================================

def compare_cities(cities: List[str]) -> dict:
    """
    比较多个城市的信息

    Args:
        cities (List[str]): 城市名称列表

    Returns:
        dict: 比较结果
    """
    city_info = {                           # 城市信息
        "北京": {"population": "2189万", "area": "16410km²"},
        "上海": {"population": "2487万", "area": "6341km²"},
        "广州": {"population": "1881万", "area": "7434km²"},
    }

    results = []                            # 结果列表
    for city in cities:                     # 遍历城市
        info = city_info.get(city)          # 获取信息
        if info:                            # 如果存在
            results.append({"city": city, **info})  # 添加到结果

    return {                                # 返回结果
        "status": "success" if results else "error",
        "results": results,
    }


# ========================================
# 示例五：健壮的工具（完善的错误处理）
# ========================================

def calculate_loan(
    principal: float,                       # 贷款本金
    annual_rate: float,                     # 年利率
    years: int,                             # 贷款年限
) -> dict:
    """
    计算贷款月供和总利息

    Args:
        principal (float): 贷款本金（元）
        annual_rate (float): 年利率（百分比，如 4.5）
        years (int): 贷款年限

    Returns:
        dict: 计算结果
    """
    # 参数验证
    if principal <= 0:                      # 本金必须大于0
        return {                            # 返回错误
            "status": "error",
            "error_code": "INVALID_PRINCIPAL",
            "error_message": "贷款本金必须大于0。"
        }

    if annual_rate <= 0 or annual_rate > 30:  # 利率范围检查
        return {                            # 返回错误
            "status": "error",
            "error_code": "INVALID_RATE",
            "error_message": "年利率必须在 0-30% 之间。"
        }

    if years <= 0 or years > 30:           # 年限范围检查
        return {                            # 返回错误
            "status": "error",
            "error_code": "INVALID_YEARS",
            "error_message": "贷款年限必须在 1-30 年之间。"
        }

    try:                                    # 尝试计算
        monthly_rate = annual_rate / 100 / 12  # 月利率
        total_months = years * 12           # 总月数

        # 等额本息月供公式
        if monthly_rate == 0:               # 如果利率为0
            monthly_payment = principal / total_months  # 直接除
        else:                               # 正常计算
            monthly_payment = (              # 月供计算
                principal * monthly_rate * (1 + monthly_rate) ** total_months
                / ((1 + monthly_rate) ** total_months - 1)
            )

        total_amount = monthly_payment * total_months  # 总还款额
        total_interest = total_amount - principal  # 总利息

        return {                            # 返回计算结果
            "status": "success",
            "monthly_payment": round(monthly_payment, 2),  # 月供
            "total_interest": round(total_interest, 2),     # 总利息
            "total_amount": round(total_amount, 2),         # 总还款额
        }

    except Exception as e:                  # 捕获计算异常
        return {                            # 返回错误
            "status": "error",
            "error_code": "CALCULATION_ERROR",
            "error_message": f"计算出错：{str(e)}"
        }


# ========================================
# 示例六：Agent-as-a-Tool
# ========================================

# 定义一个翻译 Agent
translator_agent = Agent(
    name="translator",                      # 翻译 Agent
    model="gemini-2.0-flash",               # 模型
    instruction="你是翻译专家，将文本翻译成目标语言。只输出翻译结果。",  # 指令
)

# 将翻译 Agent 包装为工具
translation_tool = AgentTool(
    agent=translator_agent,                 # 传入翻译 Agent
)


# ========================================
# 创建综合 Agent
# ========================================

comprehensive_agent = Agent(
    name="comprehensive_agent",             # Agent 名称
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 系统指令
        "你是一个多功能助手。\n"
        "你可以：\n"
        "1. 查询天气（get_weather）\n"
        "2. 搜索餐厅（search_restaurant）\n"
        "3. 创建用户档案（create_user_profile）\n"
        "4. 比较城市（compare_cities）\n"
        "5. 计算贷款（calculate_loan）\n"
        "6. 翻译文本（translation_tool）\n"
        "根据用户需求选择合适的工具。"
    ),
    tools=[                                 # 注册所有工具
        get_weather,                        # 天气查询
        search_restaurant,                  # 餐厅搜索
        create_user_profile,                # 用户档案
        compare_cities,                     # 城市比较
        calculate_loan,                     # 贷款计算
        translation_tool,                   # 翻译工具
    ],
)


# ========================================
# 运行示例
# ========================================

APP_NAME = "tools_demo"                     # 应用名称
USER_ID = "user_001"                        # 用户 ID
SESSION_ID = "session_001"                  # 会话 ID


async def main():
    """运行工具演示"""

    # 创建会话服务
    session_service = InMemorySessionService()  # 会话服务
    await session_service.create_session(    # 创建会话
        app_name=APP_NAME,                   # 应用名称
        user_id=USER_ID,                     # 用户 ID
        session_id=SESSION_ID,               # 会话 ID
    )

    # 创建运行器
    runner = Runner(
        agent=comprehensive_agent,           # Agent
        app_name=APP_NAME,                   # 应用名称
        session_service=session_service,     # 会话服务
    )

    # 测试查询
    queries = [                             # 测试问题列表
        "北京天气怎么样？",                   # 天气查询
        "帮我搜索上海的意大利餐厅",           # 餐厅搜索
        "贷款100万，利率4.5%，30年，月供多少？",  # 贷款计算
    ]

    for query in queries:                   # 遍历测试问题
        print(f"\n[用户]: {query}")         # 打印用户输入

        # 构造消息
        content = types.Content(
            role='user',                    # 角色
            parts=[types.Part(text=query)], # 内容
        )

        # 运行 Agent
        events = runner.run_async(
            user_id=USER_ID,                # 用户 ID
            session_id=SESSION_ID,          # 会话 ID
            new_message=content,            # 消息
        )

        # 获取响应
        async for event in events:           # 遍历事件
            if event.is_final_response():    # 最终响应
                print(f"[Agent]: {event.content.parts[0].text}")  # 打印回复


if __name__ == "__main__":                  # 直接运行
    asyncio.run(main())                      # 执行主函数