"""
Google ADK 会话与状态管理完整示例
展示 Session、State、Memory 的使用方法

对应教程：第06章 - 会话与状态管理
"""

# 导入 ADK 核心模块
from google.adk.agents import Agent          # Agent 类
from google.adk.runners import Runner        # 运行器
from google.adk.sessions import InMemorySessionService  # 内存会话服务
from google.adk.memory import InMemoryMemoryService  # 内存记忆服务
from google.genai import types               # 类型定义

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


# ========================================
# 示例一：State 管理（购物车）
# ========================================

def add_to_cart(item: str, price: float, ctx) -> dict:
    """
    将商品添加到购物车

    Args:
        item (str): 商品名称
        price (float): 商品价格
        ctx: 工具上下文（自动注入）

    Returns:
        dict: 操作结果
    """
    # 获取当前购物车
    cart = ctx.state.get("cart", [])         # 从 state 读取购物车

    # 添加新商品
    cart.append({                            # 追加商品
        "item": item,                        # 商品名
        "price": price,                      # 价格
    })

    # 更新 state
    ctx.state["cart"] = cart                 # 写回 state

    # 计算总价
    total = sum(item["price"] for item in cart)  # 求和

    return {                                 # 返回结果
        "status": "success",
        "message": f"已将 {item} 添加到购物车",
        "cart_items": len(cart),
        "total": total,
    }


def get_cart(ctx) -> dict:
    """
    获取购物车内容

    Args:
        ctx: 工具上下文

    Returns:
        dict: 购物车内容
    """
    cart = ctx.state.get("cart", [])         # 读取购物车

    if not cart:                             # 如果为空
        return {                             # 返回空信息
            "status": "success",
            "message": "购物车是空的",
            "items": [],
            "total": 0,
        }

    total = sum(item["price"] for item in cart)  # 计算总价

    return {                                 # 返回购物车
        "status": "success",
        "items": cart,
        "total": total,
    }


def clear_cart(ctx) -> dict:
    """
    清空购物车

    Args:
        ctx: 工具上下文

    Returns:
        dict: 操作结果
    """
    ctx.state["cart"] = []                   # 清空购物车

    return {                                 # 返回结果
        "status": "success",
        "message": "购物车已清空",
    }


# ========================================
# 示例二：temp: 临时状态
# ========================================

def step1_process(query: str, ctx) -> dict:
    """
    处理步骤一：预处理数据

    Args:
        query (str): 用户查询
        ctx: 工具上下文

    Returns:
        dict: 预处理结果
    """
    processed = query.strip().lower()       # 预处理

    # 保存到临时状态（当前调用结束后自动清除）
    ctx.state["temp:processed_query"] = processed  # 临时存储

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


def step2_enhance(ctx) -> dict:
    """
    处理步骤二：增强数据

    Args:
        ctx: 工具上下文

    Returns:
        dict: 增强结果
    """
    # 从临时状态读取
    processed = ctx.state.get("temp:processed_query")  # 读取临时状态

    if not processed:                        # 如果没有数据
        return {                             # 返回错误
            "status": "error",
            "error_message": "请先执行预处理步骤。"
        }

    enhanced = f"[增强] {processed}"         # 增强处理

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


# ========================================
# 创建 Agent
# ========================================

shopping_agent = Agent(
    name="shopping_assistant",               # Agent 名称
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令
        "你是一个购物助手。\n"
        "帮助用户管理购物车：\n"
        "- 添加商品：使用 add_to_cart\n"
        "- 查看购物车：使用 get_cart\n"
        "- 清空购物车：使用 clear_cart"
    ),
    tools=[add_to_cart, get_cart, clear_cart],  # 注册工具
)


# ========================================
# 示例三：Session 管理
# ========================================

async def session_management_demo():
    """演示 Session 管理功能"""

    print("=" * 50)                         # 分隔线
    print("Session 管理演示")                 # 标题
    print("=" * 50)                         # 分隔线

    # 创建会话服务
    session_service = InMemorySessionService()  # 内存会话服务

    # 创建会话
    session = await session_service.create_session(
        app_name="shopping_app",             # 应用名称
        user_id="user_001",                  # 用户 ID
        session_id="session_001",            # 会话 ID
        state={"cart": []},                  # 初始化状态
    )

    print(f"✅ 会话创建成功: {session.id}")   # 打印会话 ID

    # 获取会话
    existing = await session_service.get_session(
        app_name="shopping_app",
        user_id="user_001",
        session_id="session_001",
    )
    print(f"✅ 获取会话: {existing.id}")      # 打印会话信息

    # 删除会话
    await session_service.delete_session(
        app_name="shopping_app",
        user_id="user_001",
        session_id="session_001",
    )
    print("✅ 会话已删除")                    # 确认删除


# ========================================
# 示例四：多轮对话
# ========================================

async def multi_turn_demo():
    """演示多轮对话"""

    print("\n" + "=" * 50)                  # 分隔线
    print("多轮对话演示")                     # 标题
    print("=" * 50)                         # 分隔线

    # 初始化
    session_service = InMemorySessionService()  # 会话服务
    await session_service.create_session(    # 创建会话
        app_name="chat_app",                 # 应用名称
        user_id="user_001",                  # 用户 ID
        session_id="chat_001",               # 会话 ID
    )

    # 创建 Runner
    runner = Runner(
        agent=shopping_agent,                # Agent
        app_name="chat_app",                 # 应用名称
        session_service=session_service,     # 会话服务
    )

    # 模拟多轮对话
    conversations = [                       # 对话列表
        "帮我添加一个苹果，价格5元",          # 第一轮
        "再添加一个香蕉，价格3元",            # 第二轮
        "查看我的购物车",                    # 第三轮
        "清空购物车",                        # 第四轮
    ]

    for msg in conversations:               # 遍历对话
        print(f"\n[用户]: {msg}")           # 打印用户消息

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

        # 运行 Agent
        events = runner.run_async(
            user_id="user_001",              # 用户 ID
            session_id="chat_001",           # 同一会话
            new_message=content,             # 消息
        )

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


# ========================================
# 示例五：Memory 管理
# ========================================

async def memory_demo():
    """演示 Memory 管理功能"""

    print("\n" + "=" * 50)                  # 分隔线
    print("Memory 管理演示")                  # 标题
    print("=" * 50)                         # 分隔线

    # 创建记忆服务
    memory_service = InMemoryMemoryService()  # 内存记忆服务

    # 添加记忆（从会话事件中提取）
    await memory_service.add_session_events_to_memory(
        app_name="my_app",                   # 应用名称
        user_id="user_001",                  # 用户 ID
        session_id="session_001",            # 会话 ID
    )
    print("✅ 记忆已添加")                    # 确认添加

    # 搜索记忆
    results = await memory_service.search(
        query="用户偏好",                     # 搜索查询
        app_name="my_app",                   # 应用名称
        user_id="user_001",                  # 用户 ID
    )
    print(f"✅ 搜索完成，结果数: {len(results)}")  # 打印结果数


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

async def main():
    """主函数"""

    # 运行所有演示
    await session_management_demo()          # Session 管理
    await multi_turn_demo()                  # 多轮对话
    await memory_demo()                      # Memory 管理


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