"""
📌 模式3: 并行结果安全合并
核心机制: 自定义合并函数防止结果覆盖
适用场景: 多Agent协作、并行任务聚合、需要安全合并嵌套结构
"""
import asyncio
from typing import TypedDict, Dict, Annotated, Any
from langgraph.graph import StateGraph, START, END


# ===== 1. 状态定义（关键：自定义合并函数）=====
class ParallelState(TypedDict):
    """
    并行状态设计:
    - user_query: 原始查询（保留）
    - results: 使用自定义合并函数 → 安全合并嵌套字典
    - final_answer: 最终聚合结果
    """
    user_query: str
    # ⚠️ 核心：自定义合并函数实现 {**old, **new}
    results: Annotated[Dict[str, Any], lambda old, new: {**old, **new}]
    final_answer: str


# ===== 2. 并行任务节点 =====
async def task_weather(state: ParallelState) -> dict:
    """
    天气任务: 返回嵌套结果 {"weather": {...}}
    ⚠️ 必须使用任务名作为键，避免覆盖
    """
    await asyncio.sleep(0.4)  # 模拟耗时操作
    print("[🌤️ 天气任务] 完成 (0.4s)")
    return {
        "results": {
            "weather": {  # 关键：以任务名为键
                "temperature": 25,
                "condition": "sunny"
            }
        }
    }


async def task_time(state: ParallelState) -> dict:
    """
    时间任务: 返回嵌套结果 {"time": {...}}
    与天气任务互不干扰
    """
    await asyncio.sleep(0.3)  # 模拟耗时操作
    print("[⏰ 时间任务] 完成 (0.3s)")
    return {
        "results": {
            "time": {  # 关键：以任务名为键
                "hour": 15,
                "period": "afternoon"
            }
        }
    }


def aggregator_node(state: ParallelState) -> dict:
    """
    聚合节点: 综合所有并行结果
    自动等待所有触发的并行节点完成
    """
    print("\n[✅ 聚合节点] 所有并行任务完成!")
    print(f"  已收集结果: {list(state['results'].keys())}")

    # 安全访问结果（处理可能缺失的任务）
    weather = state["results"].get("weather", {})
    time = state["results"].get("time", {})

    # 生成综合决策
    decision = (
        f"📍 综合分析报告\n"
        f"{'─' * 30}\n"
        f"🌤️ 天气: {weather.get('temperature', 'N/A')}°C, {weather.get('condition', 'unknown')}\n"
        f"⏰ 时间: {time.get('hour', 'N/A')}点 ({time.get('period', 'unknown')})\n"
        f"{'─' * 30}\n"
        f"💡 建议: 下午时段天气晴朗，适合户外活动"
    )
    return {"final_answer": decision}


# ===== 3. 构建并行图 =====
def build_parallel_graph():
    builder = StateGraph(ParallelState)

    # 添加节点
    builder.add_node("weather_task", task_weather)
    builder.add_node("time_task", task_time)
    builder.add_node("aggregator", aggregator_node)

    # 并行扇出（从START同时触发两个任务）
    builder.add_edge(START, "weather_task")
    builder.add_edge(START, "time_task")

    # 扇入聚合（两个任务都完成后触发）
    builder.add_edge("weather_task", "aggregator")
    builder.add_edge("time_task", "aggregator")
    builder.add_edge("aggregator", END)

    return builder.compile()


# ===== 4. 执行演示 =====
async def main():
    print("=" * 60)
    print("🧠 模式3: 并行结果安全合并（自定义合并函数）")
    print("=" * 60)

    graph = build_parallel_graph()
    # 画图
    print(graph.get_graph().draw_ascii())

    # 初始状态（关键：results 必须初始化为空字典！）
    initial_state = {
        "user_query": "北京现在适合户外活动吗？",
        "results": {},  # ⚠️ 必须初始化，否则首次合并失败
        "final_answer": ""
    }


    # 记录时间验证并行
    start = asyncio.get_event_loop().time()
    final_state = await graph.ainvoke(initial_state)
    elapsed = asyncio.get_event_loop().time() - start

    print(f"\n【执行结果】")
    print(f"  总耗时: {elapsed:.2f}秒 ")
    print(f"\n{final_state['final_answer']}")




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

============================================================
🧠 模式3: 并行结果安全合并（自定义合并函数）
============================================================
             +-----------+                
             | __start__ |                
             +-----------+                
            ***          **               
           *               **             
         **                  **           
+-----------+           +--------------+  
| time_task |           | weather_task |  
+-----------+           +--------------+  
            ***          **               
               *       **                 
                **   **                   
            +------------+                
            | aggregator |                
            +------------+                
                   *                      
                   *                      
                   *                      
              +---------+                 
              | __end__ |                 
              +---------+                 
[⏰ 时间任务] 完成 (0.3s)
[🌤️ 天气任务] 完成 (0.4s)

[✅ 聚合节点] 所有并行任务完成!
  已收集结果: ['weather', 'time']

【执行结果】
  总耗时: 0.42秒

📍 综合分析报告
──────────────────────────────
🌤️ 天气: 25°C, sunny
⏰ 时间: 15点 (afternoon)
──────────────────────────────
💡 建议: 下午时段天气晴朗，适合户外活动