"""
📌 模式1: 简单覆盖式状态
核心机制: 默认浅合并（只更新返回字段，其他字段自动保留）
适用场景: 简单工作流、单路径执行
"""
import asyncio
from typing import TypedDict
from langgraph.graph import StateGraph, START, END


# ===== 1. 状态定义 =====
class SimpleState(TypedDict):
    """
    简单状态结构:
    - query:  用户原始查询（全程保留）
    - result: 中间/最终结果（逐步更新）
    - step:   当前执行步骤（递增）
    """
    query: str
    result: str
    step: int


# ===== 2. 节点函数 =====
def step_1_node(state: SimpleState) -> dict:
    """
    第一步: 更新 step 和部分 result
    ⚠️ 注意: 未返回 'query' 字段，但会被自动保留！
    """
    print(f"[Step1] 输入状态: {state}")
    return {
        "step": state["step"] + 1,  # 更新 step
        "result": "初步分析完成"  # 更新 result
        # query 字段未返回 → 自动保留原始值
    }


def step_2_node(state: SimpleState) -> dict:
    """
    第二步: 只更新 result，step 保持不变
    """
    print(f"[Step2] 输入状态: {state}")
    return {
        "result": f"{state['result']} → 最终结论生成"  # 仅更新 result
        # step 和 query 自动保留
    }


# ===== 3. 构建图 =====
def build_simple_graph():
    builder = StateGraph(SimpleState)
    builder.add_node("step_1", step_1_node)
    builder.add_node("step_2", step_2_node)

    builder.add_edge(START, "step_1")
    builder.add_edge("step_1", "step_2")
    builder.add_edge("step_2", END)

    return builder.compile()


# ===== 4. 执行演示 =====
async def main():
    print("=" * 60)
    print("🧠 模式1: 简单覆盖式状态（默认浅合并）")
    print("=" * 60)

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

    # 初始状态（包含所有字段）
    initial_state = {
        "query": "北京天气如何？",
        "result": "初始状态",
        "step": 0
    }

    print("\n【初始状态】")
    print(f"  {initial_state}")

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

    print("\n【最终状态】")
    print(f"  {final_state}")




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