"""LangGraph 并行节点与 Reducer 示例。

本例包含两个图：

1. good_graph：两个搜索节点并行执行，并用 reducer 合并结果。
2. bad_graph：两个并行节点同时覆盖同一字段，触发 InvalidUpdateError。

主图结构：

                       ┌─> search_web ─────┐
    START -> prepare ──┤                   ├─> summarize -> END
                       └─> search_database ┘

运行：
    python 10_parallel_and_reducers.py
"""

import operator
import time
from typing import Annotated, TypedDict

from langgraph.errors import InvalidUpdateError
from langgraph.graph import END, START, StateGraph


# -----------------------------------------------------------------------------
# 1. Reducer
# -----------------------------------------------------------------------------


def merge_unique(left: list[str], right: list[str]) -> list[str]:
    """合并两个列表并去重，同时保留第一次出现的顺序。"""
    return list(dict.fromkeys(left + right))


class ResearchState(TypedDict):
    topic: str

    # operator.add：把各节点返回的列表连接起来。
    results: Annotated[list[str], operator.add]

    # 自定义 reducer：合并日志并去除重复内容。
    logs: Annotated[list[str], merge_unique]

    # 没有 reducer：使用默认覆盖策略。
    # 这里只有 summarize 节点更新它，因此不会发生并行冲突。
    summary: str


# -----------------------------------------------------------------------------
# 2. 主图节点
# -----------------------------------------------------------------------------


def prepare(state: ResearchState) -> dict:
    """准备搜索；完成后两个下游节点会在同一 super-step 并行执行。"""
    print(f"[prepare] 准备研究：{state['topic']}")
    return {"logs": ["准备完成"]}


def search_web(state: ResearchState) -> dict:
    """模拟耗时 1 秒的网页搜索。"""
    print("[search_web] 开始")
    time.sleep(1)
    print("[search_web] 完成")
    return {
        "results": [
            f"网页结果：{state['topic']} 支持有状态工作流",
            "共同结果：LangGraph 使用 StateGraph 编排节点",
        ],
        "logs": ["网页搜索完成", "搜索阶段完成"],
    }


def search_database(state: ResearchState) -> dict:
    """模拟耗时 1 秒的数据库搜索。"""
    print("[search_database] 开始")
    time.sleep(1)
    print("[search_database] 完成")
    return {
        "results": [
            f"数据库结果：{state['topic']} 支持持久化和恢复",
            "共同结果：LangGraph 使用 StateGraph 编排节点",
        ],
        "logs": ["数据库搜索完成", "搜索阶段完成"],
    }


def summarize(state: ResearchState) -> dict:
    """汇合节点：两个搜索分支都结束后才会执行。"""
    print(f"[summarize] 收到 {len(state['results'])} 条结果")

    lines = "\n".join(
        f"{index}. {result}"
        for index, result in enumerate(state["results"], start=1)
    )
    return {
        "summary": f"主题：{state['topic']}\n{lines}",
        "logs": ["摘要生成完成"],
    }


# -----------------------------------------------------------------------------
# 3. 构建正确的并行图
# -----------------------------------------------------------------------------


good_builder = StateGraph(ResearchState)
good_builder.add_node("prepare", prepare)
good_builder.add_node("search_web", search_web)
good_builder.add_node("search_database", search_database)
good_builder.add_node("summarize", summarize)

good_builder.add_edge(START, "prepare")

# 一个节点拥有两条普通出边：下一 super-step 中两个目标节点并行执行。
good_builder.add_edge("prepare", "search_web")
good_builder.add_edge("prepare", "search_database")

# 两个分支都指向 summarize。LangGraph 会等待同一轮的两个分支完成后汇合。
good_builder.add_edge("search_web", "summarize")
good_builder.add_edge("search_database", "summarize")
good_builder.add_edge("summarize", END)

good_graph = good_builder.compile()


# -----------------------------------------------------------------------------
# 4. 错误示例：并行节点覆盖同一个无 reducer 字段
# -----------------------------------------------------------------------------


class BadState(TypedDict):
    value: str  # 没有 reducer，默认策略是覆盖


def left_branch(state: BadState) -> dict:
    return {"value": "来自左分支"}


def right_branch(state: BadState) -> dict:
    return {"value": "来自右分支"}


bad_builder = StateGraph(BadState)
bad_builder.add_node("left", left_branch)
bad_builder.add_node("right", right_branch)
bad_builder.add_edge(START, "left")
bad_builder.add_edge(START, "right")
bad_builder.add_edge("left", END)
bad_builder.add_edge("right", END)
bad_graph = bad_builder.compile()


# -----------------------------------------------------------------------------
# 5. 演示函数
# -----------------------------------------------------------------------------


def demo_good_graph() -> None:
    print("=== 正确示例：并行 + Reducer + 汇合 ===")

    initial_state: ResearchState = {
        "topic": "LangGraph",
        "results": [],
        "logs": [],
        "summary": "",
    }

    started = time.perf_counter()
    result = good_graph.invoke(initial_state)
    elapsed = time.perf_counter() - started

    print(f"\n耗时：{elapsed:.2f} 秒")
    print("两个搜索节点各等待 1 秒；总耗时约 1 秒，说明它们并行执行。")

    print("\nresults（operator.add，保留重复项）：")
    for item in result["results"]:
        print(f"- {item}")

    print("\nlogs（merge_unique，重复日志只保留一次）：")
    for item in result["logs"]:
        print(f"- {item}")

    print("\nsummary（默认覆盖，由 summarize 单独写入）：")
    print(result["summary"])


def demo_stream_updates() -> None:
    print("\n=== 使用 updates 观察并行 super-step ===")

    initial_state: ResearchState = {
        "topic": "LangGraph",
        "results": [],
        "logs": [],
        "summary": "",
    }

    for part in good_graph.stream(
        initial_state,
        stream_mode="updates",
        version="v2",
    ):
        if part["type"] != "updates":
            continue

        node_names = ", ".join(part["data"].keys())
        print(f"节点更新：{node_names}")


def demo_conflict() -> None:
    print("\n=== 错误示例：并行更新冲突 ===")
    try:
        bad_graph.invoke({"value": "初始值"})
    except InvalidUpdateError as exc:
        print(f"捕获到：{type(exc).__name__}")
        print("原因：left 和 right 在同一个 super-step 中都写入 value，")
        print("而 value 没有 reducer，LangGraph 不知道应该保留哪一个值。")
        print(f"原始错误：{exc}")


def main() -> None:
    demo_good_graph()
    demo_stream_updates()
    demo_conflict()

    print("\n=== 学习结论 ===")
    print("1. 多条普通出边会让目标节点在同一个 super-step 中并行运行。")
    print("2. 并行分支汇合时，下游节点等待本轮所有上游分支完成。")
    print("3. 无 reducer 的字段采用覆盖策略，同一步只能接收一个更新。")
    print("4. operator.add 可以连接列表，但不会去重。")
    print("5. 自定义 reducer 可以实现去重、排序或其他合并规则。")


if __name__ == "__main__":
    main()