"""
并行 + reducer 追踪版：把 menu 通道的每一次变化打出来。

04_langgraph_paramlizm_demo.py 是"错误写法"的现场，本文件是它的显微镜：
把 operator.add 换成一个会打印的 traced_add，就能看见 menu 每一步
从什么变成什么，从而搞清楚 3 个菜为什么会变成 24 个。

用法：
    python -X utf8 05_langgraph_paramlizm_demo2.py            # 错误/正确 两种写法对比跑
    python -X utf8 05_langgraph_paramlizm_demo2.py --mode wrong
    python -X utf8 05_langgraph_paramlizm_demo2.py --mode right
    python -X utf8 05_langgraph_paramlizm_demo2.py --stream    # 再用 stream 看增量/快照
    python -X utf8 05_langgraph_paramlizm_demo2.py --no-sleep  # 去掉 sleep，跑得快点

配套文档：04_并行与reducer说明.md
"""

import argparse
import time
from operator import add
from typing import Annotated, List, TypedDict

from langgraph.graph import END, START, StateGraph

# 是否真的 sleep（--no-sleep 关掉）
SLEEP = True
# reducer 被调用的次数，用来给日志编号
_reduce_calls = 0


def cook(seconds: float) -> None:
    """模拟做菜耗时。"""
    if SLEEP:
        time.sleep(seconds)


def traced_add(old: List[str], new: List[str]) -> List[str]:
    """带日志的 reducer。

    reducer 的签名永远是 (通道现在的值, 节点返回的值) -> 通道的新值。
    这里除了照常拼接，还把两边的内容和长度打出来 —— 这就是"追踪 menu"的关键。
    注意 reducer 是**每个节点的更新来一次**，不是一次收到所有并行结果。
    """
    global _reduce_calls
    _reduce_calls += 1
    result = add(old, new)
    print(
        f"   🔧 [reducer #{_reduce_calls}] "
        f"old({len(old)})={old}  +  new({len(new)})={new}  ->  长度 {len(result)}"
    )
    return result


class MyState(TypedDict):
    menu: Annotated[List[str], traced_add]  # 有 reducer：节点返回值是"增量"，会被追加
    all_ready: bool  # 无 reducer：节点返回值是"新值"，直接覆盖


def peek(node: str, state: MyState) -> None:
    """打印节点看到的 menu 快照，以及那个 list 对象的 id。

    并行节点打出来的 id 是同一个 —— 这说明 state 里的 list 是共享对象，
    所以 append 会直接改到通道里的数据上，绕过 reducer。
    """
    menu = state.get("menu", [])
    print(f"   👀 {node} 看到 menu({len(menu)}) = {menu}   id={id(menu)}")


# ======================================================================
# 写法一：错误写法（04 里的那份）—— append 改共享 list + return state
# ======================================================================
def wrong_start(state: MyState):
    print("👨‍🍳 [错误写法] 厨房准备开始！")
    peek("start", state)
    state["menu"] = []  # 多余：带 reducer 的通道本来就是 []
    return state


def wrong_chef_a(state: MyState):
    print("👩‍🍳 A 厨师在做：番茄炒蛋...")
    peek("chef_a", state)
    cook(2)
    state["menu"].append("番茄炒蛋")  # ① 就地改了三个节点共享的那个 list
    return state  # ② 又把改完的整份当增量交回去


def wrong_chef_b(state: MyState):
    print("👨‍🍳 B 厨师在做：宫保鸡丁...")
    peek("chef_b", state)
    cook(3)
    state["menu"].append("宫保鸡丁")
    return state


def wrong_chef_c(state: MyState):
    print("👩‍🍳 C 厨师在做：鱼香肉丝...")
    peek("chef_c", state)
    cook(1)
    state["menu"].append("鱼香肉丝")
    return state


def wrong_serve(state: MyState):
    print("\n🍽️ [错误写法] 上菜：")
    peek("serve", state)
    for dish in state["menu"]:
        print("   -", dish)
    state["all_ready"] = True
    return state  # 汇聚节点的 return state：把读到的整份 menu 再追加一遍


# ======================================================================
# 写法二：正确写法 —— 只返回自己新增的那一份，不碰 state 里的 list
# ======================================================================
def right_start(state: MyState):
    print("👨‍🍳 [正确写法] 厨房准备开始！")
    peek("start", state)
    return {}  # 什么都不写就返回空 dict


def right_chef_a(state: MyState):
    print("👩‍🍳 A 厨师在做：番茄炒蛋...")
    peek("chef_a", state)
    cook(2)
    return {"menu": ["番茄炒蛋"]}  # 只交增量，合并交给 reducer


def right_chef_b(state: MyState):
    print("👨‍🍳 B 厨师在做：宫保鸡丁...")
    peek("chef_b", state)
    cook(3)
    return {"menu": ["宫保鸡丁"]}


def right_chef_c(state: MyState):
    print("👩‍🍳 C 厨师在做：鱼香肉丝...")
    peek("chef_c", state)
    cook(1)
    return {"menu": ["鱼香肉丝"]}


def right_serve(state: MyState):
    print("\n🍽️ [正确写法] 上菜：")
    peek("serve", state)
    for dish in state["menu"]:
        print("   -", dish)
    return {"all_ready": True}  # 绝不带 menu，否则整份再追加一遍


# --- 构建 LangGraph（两种写法共用同一张图，只换节点函数）---
NODES = {
    "wrong": (wrong_start, wrong_chef_a, wrong_chef_b, wrong_chef_c, wrong_serve),
    "right": (right_start, right_chef_a, right_chef_b, right_chef_c, right_serve),
}


def build_graph(mode: str):
    start, a, b, c, serve = NODES[mode]

    graph_builder = StateGraph(MyState)
    graph_builder.add_node("start", start)
    graph_builder.add_node("chef_a", a)
    graph_builder.add_node("chef_b", b)
    graph_builder.add_node("chef_c", c)
    graph_builder.add_node("serve", serve)

    # 启动 → 并行三个厨师
    graph_builder.add_edge(START, "start")
    graph_builder.add_edge("start", "chef_a")
    graph_builder.add_edge("start", "chef_b")
    graph_builder.add_edge("start", "chef_c")

    # 三个厨师汇总到 serve（fan-in）
    graph_builder.add_edge(["chef_a", "chef_b", "chef_c"], "serve")
    graph_builder.add_edge("serve", END)

    return graph_builder.compile()


# --- 两种查看方式 ---
def run_invoke(mode: str) -> List[str]:
    """普通 invoke：靠 reducer 日志看 menu 的每一次合并。"""
    global _reduce_calls
    _reduce_calls = 0

    title = "错误写法：append + return state" if mode == "wrong" else "正确写法：只返回增量"
    print("\n" + "=" * 70)
    print(f"  {title}   (mode={mode})")
    print("=" * 70)

    result = build_graph(mode).invoke({})
    menu = result["menu"]
    print(f"\n📊 reducer 共被调用 {_reduce_calls} 次，最终 menu 长度 = {len(menu)}（期望 3）")
    print(f"   最终结果: {result}")
    return menu


def run_stream(mode: str) -> None:
    """stream 方式：把"节点交了什么"和"通道变成了什么"分开看。

    stream_mode="updates"（默认）：每个节点跑完 yield 一次 {节点名: 该节点的返回值}。
        看到的是**增量**，也就是交给 reducer 的那个 new。
        错误写法在这里最扎眼 —— chef_a 交回来的不是 1 个菜而是一整份。
    stream_mode="values"：通道值变化后 yield 一次完整快照。
        看到的是**合并后的结果**，也就是 reducer 返回的 old。
        并行的三个厨师合并完才给一份，所以快照数比节点数少。
    """
    global _reduce_calls

    print("\n" + "=" * 70)
    print(f"  stream_mode='updates'：每个节点交回的增量 (mode={mode})")
    print("=" * 70)
    _reduce_calls = 0
    for i, chunk in enumerate(build_graph(mode).stream({}), start=1):
        for node_name, update in chunk.items():
            menu = update.get("menu") if isinstance(update, dict) else None
            n = len(menu) if isinstance(menu, list) else "-"
            print(f"\n── 更新 {i} ── {node_name} 交回 menu({n}) = {menu}")

    print("\n" + "=" * 70)
    print(f"  stream_mode='values'：合并后的通道快照 (mode={mode})")
    print("=" * 70)
    _reduce_calls = 0
    for i, snap in enumerate(build_graph(mode).stream({}, stream_mode="values")):
        menu = snap.get("menu", [])
        print(f"\n── 快照 {i} ── menu({len(menu)}) = {menu}  all_ready={snap.get('all_ready')}")


# --- 运行 ---
def main() -> None:
    global SLEEP

    ap = argparse.ArgumentParser(description="LangGraph 并行 + reducer 追踪 demo")
    ap.add_argument(
        "--mode",
        choices=["wrong", "right", "both"],
        default="both",
        help="跑哪种写法，默认两种都跑做对比",
    )
    ap.add_argument("--stream", action="store_true", help="额外用 stream 看增量和快照")
    ap.add_argument("--no-sleep", action="store_true", help="去掉模拟耗时，跑得快点")
    args = ap.parse_args()

    SLEEP = not args.no_sleep
    modes = ["wrong", "right"] if args.mode == "both" else [args.mode]

    menus = {m: run_invoke(m) for m in modes}

    if args.stream:
        for m in modes:
            run_stream(m)

    if args.mode == "both":
        print("\n" + "=" * 70)
        print("  对比")
        print("=" * 70)
        print(f"  错误写法 menu 长度 = {len(menus['wrong'])}   (3 个菜被写了 8 遍)")
        print(f"  正确写法 menu 长度 = {len(menus['right'])}   ✅")
        print("\n  两级翻倍：")
        print("    3 → 12  三次 append 改共享 list，又被三次 return state 各追加 3 条")
        print("    12 → 24 serve 的 return state 把读到的 12 条整份再追加一遍")
        print("\n  正确写法的 menu 顺序是 a/b/c 的声明顺序，不是完成时间顺序")
        print("  （c 只睡 1 秒却排最后）—— add 按 LangGraph 处理更新的顺序拼接。")


if __name__ == "__main__":
    main()