"""
Google ADK 多智能体系统完整示例
展示各种多 Agent 协作模式

对应教程：第05章 - 多智能体系统
"""

# 导入 ADK 核心模块
from google.adk.agents import (                # 导入所有 Agent 类型
    Agent,                                     # LLM Agent（别名）
    LlmAgent,                                  # LLM Agent（完整名）
    SequentialAgent,                           # 顺序工作流
    ParallelAgent,                             # 并行工作流
    LoopAgent,                                 # 循环工作流
    BaseAgent,                                 # 基础 Agent（自定义用）
)
from google.adk.events import Event, EventActions  # 导入事件类型
from google.adk.agents.invocation_context import InvocationContext  # 导入上下文
from typing import AsyncGenerator              # 导入异步生成器


# ========================================
# 模式一：协调器/调度器模式
# ========================================

# 定义专门的子 Agent
math_agent = LlmAgent(
    name="MathExpert",                      # 数学专家
    model="gemini-2.0-flash",               # 模型
    description="擅长数学计算和统计分析。",  # 描述
    instruction="你是数学专家，解决数学问题。给出详细的计算过程。",  # 指令
)

writing_agent = LlmAgent(
    name="WritingExpert",                   # 写作专家
    model="gemini-2.0-flash",               # 模型
    description="擅长文案撰写和内容创作。",  # 描述
    instruction="你是写作专家，帮助用户撰写高质量的内容。",  # 指令
)

coding_agent = LlmAgent(
    name="CodingExpert",                    # 编程专家
    model="gemini-2.0-flash",               # 模型
    description="擅长编程和代码调试。",      # 描述
    instruction="你是编程专家，帮助用户解决编程问题。提供带注释的代码。",  # 指令
)

# 定义协调器
coordinator = LlmAgent(
    name="Coordinator",                     # 协调器
    model="gemini-2.0-flash",               # 模型
    description="根据问题类型分配给合适的专家。",  # 描述
    instruction=(                           # 指令
        "你是一个任务协调器。\n"
        "根据用户的问题类型，将任务分配给合适的专家：\n"
        "- 数学相关 → MathExpert\n"
        "- 写作相关 → WritingExpert\n"
        "- 编程相关 → CodingExpert\n"
        "使用 transfer_to_agent 工具进行任务分配。"
    ),
    sub_agents=[math_agent, writing_agent, coding_agent],  # 注册子 Agent
)


# ========================================
# 模式二：顺序流水线模式
# ========================================

# 步骤一：信息提取
extractor = LlmAgent(
    name="Extractor",                       # 信息提取 Agent
    model="gemini-2.0-flash",               # 模型
    instruction="从用户输入中提取关键信息。只输出提取的关键信息。",  # 指令
    output_key="extracted_info",            # 输出保存到 state
)

# 步骤二：信息分析
analyzer = LlmAgent(
    name="Analyzer",                        # 信息分析 Agent
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令：引用上一步输出
        "分析以下提取的信息，给出深度分析报告。\n"
        "提取的信息：{extracted_info}"
    ),
    output_key="analysis_result",           # 输出保存到 state
)

# 步骤三：报告生成
reporter = LlmAgent(
    name="Reporter",                        # 报告生成 Agent
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令：引用前两步输出
        "根据分析结果生成一份简洁的报告。\n"
        "原始信息：{extracted_info}\n"
        "分析结果：{analysis_result}"
    ),
)

# 创建顺序流水线
pipeline = SequentialAgent(
    name="InfoPipeline",                    # 流水线名称
    sub_agents=[extractor, analyzer, reporter],  # 按顺序执行
)


# ========================================
# 模式三：并行扇出模式
# ========================================

# 并行任务一：天气查询
weather_agent = LlmAgent(
    name="WeatherFetcher",                  # 天气查询
    model="gemini-2.0-flash",               # 模型
    instruction="查询并返回天气信息。只返回天气数据。",  # 指令
    output_key="weather_data",              # 输出保存到 state
)

# 并行任务二：新闻查询
news_agent = LlmAgent(
    name="NewsFetcher",                     # 新闻查询
    model="gemini-2.0-flash",               # 模型
    instruction="查询并返回最新新闻。只返回新闻摘要。",  # 指令
    output_key="news_data",                 # 输出保存到 state
)

# 并行任务三：交通查询
traffic_agent = LlmAgent(
    name="TrafficFetcher",                  # 交通查询
    model="gemini-2.0-flash",               # 模型
    instruction="查询并返回交通状况。只返回交通信息。",  # 指令
    output_key="traffic_data",              # 输出保存到 state
)

# 创建并行执行器
gatherer = ParallelAgent(
    name="InfoGatherer",                    # 并行执行器名称
    sub_agents=[weather_agent, news_agent, traffic_agent],  # 并行执行
)


# ========================================
# 模式四：循环优化模式（Generator-Critic）
# ========================================

# 生成器 Agent
generator = LlmAgent(
    name="Generator",                       # 生成器
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令
        "你是一个创意写手。\n"
        "根据用户需求创作内容。\n"
        "将创作的内容保存到 state['draft']。"
    ),
    output_key="draft",                     # 输出保存到 state
)

# 评审器 Agent
critic = LlmAgent(
    name="Critic",                          # 评审器
    model="gemini-2.0-flash",               # 模型
    instruction=(                           # 指令
        "你是一个严格的内容评审。\n"
        "评审以下内容，给出评分（1-10）和改进建议。\n"
        "当前草稿：{draft}\n"
        "如果评分 >= 8，将 state['approved'] 设为 True。\n"
        "否则，将改进建议保存到 state['feedback']。"
    ),
)

# 条件检查 Agent
class ApprovalChecker(BaseAgent):
    """检查内容是否通过审批"""

    async def _run_async_impl(
        self,
        ctx: InvocationContext,              # 调用上下文
    ) -> AsyncGenerator[Event, None]:        # 返回事件生成器
        """执行检查"""

        approved = ctx.session.state.get("approved", False)  # 获取审批状态

        yield Event(                         # 生成事件
            author=self.name,                # 事件作者
            actions=EventActions(
                escalate=approved,           # 通过则退出循环
            ),
        )

# 创建迭代优化循环
review_loop = LoopAgent(
    name="ReviewLoop",                      # 循环名称
    max_iterations=3,                       # 最多迭代3次
    sub_agents=[critic, ApprovalChecker(name="Checker"), generator],  # 循环体
)


# ========================================
# 模式五：多层级嵌套
# ========================================

# 第三层：具体执行 Agent
code_reviewer = LlmAgent(
    name="CodeReviewer",                    # 代码审查
    model="gemini-2.0-flash",               # 模型
    description="审查代码质量。",            # 描述
    instruction="审查代码并给出改进建议。",  # 指令
)

test_writer = LlmAgent(
    name="TestWriter",                      # 测试编写
    model="gemini-2.0-flash",               # 模型
    description="编写单元测试。",            # 描述
    instruction="为代码编写单元测试。",      # 指令
)

# 第二层：开发流水线
dev_pipeline = SequentialAgent(
    name="DevPipeline",                     # 流水线
    sub_agents=[code_reviewer, test_writer], # 顺序执行
)

# 第二层：其他 Agent
doc_writer = LlmAgent(
    name="DocWriter",                       # 文档编写
    model="gemini-2.0-flash",               # 模型
    description="编写技术文档。",            # 描述
    instruction="编写清晰的技术文档。",      # 指令
)

# 第一层：根协调器
root_agent = LlmAgent(
    name="RootCoordinator",                 # 根协调器
    model="gemini-2.0-flash",               # 模型
    description="开发团队协调器。",          # 描述
    instruction=(                           # 指令
        "你是一个开发团队协调器。\n"
        "代码审查和测试交给 DevPipeline。\n"
        "文档编写交给 DocWriter。"
    ),
    sub_agents=[dev_pipeline, doc_writer],  # 子 Agent
)


# ========================================
# 打印结构信息
# ========================================

if __name__ == "__main__":
    print("=" * 60)                         # 分隔线
    print("Google ADK 多智能体系统示例")       # 标题
    print("=" * 60)                         # 分隔线

    print("\n📦 模式一：协调器/调度器")       # 模式一
    print(f"   根 Agent: {coordinator.name}")  # 根 Agent
    print(f"   子 Agent: {[a.name for a in coordinator.sub_agents]}")  # 子 Agent

    print("\n📦 模式二：顺序流水线")           # 模式二
    print(f"   流水线: {pipeline.name}")      # 流水线
    print(f"   步骤: {[a.name for a in pipeline.sub_agents]}")  # 步骤

    print("\n📦 模式三：并行扇出")             # 模式三
    print(f"   并行器: {gatherer.name}")      # 并行器
    print(f"   任务: {[a.name for a in gatherer.sub_agents]}")  # 任务

    print("\n📦 模式四：循环优化")             # 模式四
    print(f"   循环: {review_loop.name}")     # 循环
    print(f"   最大迭代: {review_loop.max_iterations}")  # 最大迭代

    print("\n📦 模式五：多层级嵌套")           # 模式五
    print(f"   根: {root_agent.name}")        # 根
    print(f"   第二层: {[a.name for a in root_agent.sub_agents]}")  # 第二层
    print(f"   第三层: {[a.name for a in dev_pipeline.sub_agents]}")  # 第三层

    print("\n✅ 所有模式定义完成！")           # 成功信息
    print("使用 'adk web' 或 'adk run' 运行。")  # 运行提示