"""LangGraph 容错示例：重试、错误恢复与循环保护。

本文件不调用真实 LLM/API，可以直接运行并稳定复现实验：

1. transient：前两次 ConnectionError，第三次成功
2. permanent：ConnectionError 重试耗尽后进入 error_handler
3. invalid：ValueError 不重试，直接进入 error_handler
4. loop：故意无限循环，由 recursion_limit 终止

需要 langgraph >= 1.2（节点级 error_handler / NodeError 是 1.2+ 能力）。

运行：
    python 09_retry_and_errors.py
"""

from typing import Literal, TypedDict

from langgraph.errors import GraphRecursionError, NodeError
from langgraph.graph import END, START, StateGraph
from langgraph.runtime import Runtime
from langgraph.types import Command, RetryPolicy


# -----------------------------------------------------------------------------
# 1. 节点重试 + 重试耗尽后的错误恢复
# -----------------------------------------------------------------------------


class ApiState(TypedDict):
    mode: Literal["transient", "permanent", "invalid"]
    result: str
    status: str
    error: str


def call_unstable_api(
    state: ApiState,
    runtime: Runtime,
) -> Command[Literal["finalize"]]:
    """模拟一个不稳定的外部 API。

    注意：失败尝试产生的 State 更新不会提交，因此不要用 State 字段记录
    当前重试次数。LangGraph 已通过 execution_info 提供 node_attempt。
    """
    attempt = runtime.execution_info.node_attempt
    mode = state["mode"]
    print(f"  call_unstable_api：第 {attempt} 次尝试，mode={mode}")

    if mode == "transient" and attempt < 3:
        raise ConnectionError("模拟网络抖动")

    if mode == "permanent":
        raise ConnectionError("模拟上游服务持续不可用")

    if mode == "invalid":
        # ValueError 通常代表参数/业务错误，重复请求不会自行恢复。
        raise ValueError("模拟无效请求参数")

    return Command(
        update={
            "result": "API 调用成功",
            "status": "success",
            "error": "",
        },
        goto="finalize",
    )


def api_error_handler(
    state: ApiState,
    error: NodeError,
) -> Command[Literal["finalize"]]:
    """在不重试或重试耗尽后，将异常转成可控的业务结果。"""
    print(
        f"  error_handler：节点 {error.node} 最终失败："
        f"{type(error.error).__name__}: {error.error}"
    )
    return Command(
        update={
            "result": "已启用降级结果，请稍后重试",
            "status": "recovered",
            "error": f"{type(error.error).__name__}: {error.error}",
        },
        goto="finalize",
    )


def finalize(state: ApiState) -> dict:
    """统一收尾节点；成功和降级路径都会到达这里。"""
    print(f"  finalize：status={state['status']}, result={state['result']}")
    return {}


api_builder = StateGraph(ApiState)
api_builder.add_node(
    "call_api",
    call_unstable_api,
    retry_policy=RetryPolicy(
        max_attempts=3,          # 包含第一次，共最多执行 3 次
        initial_interval=0.1,   # 第一次重试前等待 0.1 秒
        backoff_factor=2.0,     # 后续等待时间指数增长
        max_interval=1.0,
        jitter=False,           # 教学输出保持可预测
        retry_on=ConnectionError,
    ),
    error_handler=api_error_handler,
)
api_builder.add_node("finalize", finalize)
api_builder.add_edge(START, "call_api")
api_builder.add_edge("finalize", END)
api_graph = api_builder.compile()


def run_api_case(mode: Literal["transient", "permanent", "invalid"]) -> None:
    """运行一个重试/错误恢复场景，并打印最终状态。"""
    print(f"\n=== 场景：{mode} ===")
    result = api_graph.invoke(
        {
            "mode": mode,
            "result": "",
            "status": "pending",
            "error": "",
        }
    )
    print("  最终状态：", result)


# -----------------------------------------------------------------------------
# 2. 循环保护：recursion_limit
# -----------------------------------------------------------------------------


class LoopState(TypedDict):
    count: int


def loop_forever(state: LoopState) -> dict:
    """故意不提供终止条件，用于演示 GraphRecursionError。"""
    next_count = state["count"] + 1
    print(f"  loop：count={next_count}")
    return {"count": next_count}


loop_builder = StateGraph(LoopState)
loop_builder.add_node("loop", loop_forever)
loop_builder.add_edge(START, "loop")
loop_builder.add_edge("loop", "loop")
loop_graph = loop_builder.compile()


def run_loop_case() -> None:
    """用较小的递归限制安全终止错误设计的循环图。"""
    print("\n=== 场景：loop / recursion_limit ===")
    try:
        loop_graph.invoke(
            {"count": 0},
            config={"recursion_limit": 4},
        )
    except GraphRecursionError as exc:
        print(f"  已捕获 {type(exc).__name__}")
        print("  原因：图在限制步数内没有到达 END。")
        print("  正确做法：为循环增加业务终止条件；recursion_limit 只是最后防线。")


# -----------------------------------------------------------------------------
# 3. 入口
# -----------------------------------------------------------------------------


def main() -> None:
    print("=== LangGraph Retry & Errors ===")
    run_api_case("transient")
    run_api_case("permanent")
    run_api_case("invalid")
    run_loop_case()

    print("\n=== 观察结论 ===")
    print("1. ConnectionError 会按照 RetryPolicy 自动重试。")
    print("2. ValueError 不符合 retry_on，第一次失败后立即交给 error_handler。")
    print("3. 重试耗尽后，error_handler 可用 Command 更新状态并走降级路径。")
    print("4. recursion_limit 能阻止失控循环，但不能替代正确的终止条件。")


if __name__ == "__main__":
    main()