"""
📌 模式2: 消息历史追加
核心机制: Annotated[List, add] 实现自动追加（非覆盖）
适用场景: 聊天机器人、对话系统、需要保留完整历史的场景
"""
import asyncio
from typing import TypedDict, List, Annotated
from operator import add  # 关键：列表追加操作符
from langgraph.graph import StateGraph, START, END
from langchain_core.messages import HumanMessage, AIMessage


# ===== 1. 状态定义（关键：Annotated + add）=====
class ChatState(TypedDict):
    """
    聊天状态设计:
    - messages: 使用 Annotated[List, add] → 自动追加新消息
    - current_topic: 当前对话主题（普通字段，覆盖更新）
    """
    messages: Annotated[List[HumanMessage | AIMessage], add]  # ⚠️ 核心：add 操作符
    current_topic: str


# ===== 2. 节点函数 =====
def user_input_node(state: ChatState) -> dict:
    """
    用户输入节点: 追加用户消息
    返回 {"messages": [新消息]} → 自动追加到历史
    """
    print(f"\n[用户输入] 当前消息数: {len(state['messages'])}")
    new_msg = HumanMessage(content="用户: 今天天气怎么样？")
    return {
        "messages": [new_msg],  # 返回单元素列表
        "current_topic": "天气查询"
    }


def ai_response_node(state: ChatState) -> dict:
    """
    AI回复节点: 追加AI消息
    ⚠️ 注意: 返回的 messages 会追加，不会覆盖历史！
    """
    print(f"[AI回复] 当前消息数: {len(state['messages'])}")
    new_msg = AIMessage(content="AI: 今天北京晴，25°C，适合外出")
    return {
        "messages": [new_msg],  # 再次返回单元素列表
        "current_topic": "天气详情"
    }


def summary_node(state: ChatState) -> dict:
    """
    总结节点: 展示完整对话历史
    """
    print(f"\n[对话总结] 总消息数: {len(state['messages'])}")
    for i, msg in enumerate(state['messages'], 1):
        role = "👤" if isinstance(msg, HumanMessage) else "🤖"
        print(f"  {role} 消息{i}: {msg.content[:30]}...")
    return {"current_topic": "对话结束"}


# ===== 3. 构建图 =====
def build_chat_graph():
    builder = StateGraph(ChatState)
    builder.add_node("user_input", user_input_node)
    builder.add_node("ai_response", ai_response_node)
    builder.add_node("summary", summary_node)

    builder.add_edge(START, "user_input")
    builder.add_edge("user_input", "ai_response")
    builder.add_edge("ai_response", "summary")
    builder.add_edge("summary", END)

    return builder.compile()


# ===== 4. 执行演示 =====
async def main():
    print("=" * 60)
    print("🧠 模式2: 消息历史追加（Annotated + add）")
    print("=" * 60)

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

    # 初始状态（必须初始化 messages 为空列表！）
    initial_state = {
        "messages": [],  # ⚠️ 关键：必须初始化，否则合并失败
        "current_topic": "初始"
    }

    print("\n【初始状态】messages 长度: 0")

    # 执行图
    final_state = await graph.ainvoke(initial_state)

    print("\n【最终状态】")
    print(f"  messages 长度: {len(final_state['messages'])}")
    print(f"  current_topic: {final_state['current_topic']}")




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

============================================================
🧠 模式2: 消息历史追加（Annotated + add）
============================================================
 +-----------+   
 | __start__ |   
 +-----------+   
        *        
        *        
        *        
+------------+   
| user_input |   
+------------+   
        *        
        *        
        *        
+-------------+  
| ai_response |  
+-------------+  
        *        
        *        
        *        
  +---------+    
  | summary |    
  +---------+    
        *        
        *        
        *        
  +---------+    
  | __end__ |    
  +---------+

【初始状态】messages 长度: 0

[用户输入] 当前消息数: 0
[AI回复] 当前消息数: 1

[对话总结] 总消息数: 2
  👤 消息1: 用户: 今天天气怎么样？...
  🤖 消息2: AI: 今天北京晴，25°C，适合外出...

【最终状态】
  messages 长度: 2
  current_topic: 对话结束