"""通过 LangGraph custom 流发送进度，不额外保存任务状态"""

import asyncio
from contextlib import contextmanager
from time import perf_counter
from uuid import uuid4

from langgraph.config import get_config, get_stream_writer
from langgraph.errors import GraphInterrupt

from app.schemas.node_progress import NodeProgress, NodeProgressStatus

# 将内部节点名映射为用户能理解的进度说明，未配置的节点仍发送事件但不附带 label
NODE_PROGRESS_MESSAGES = {
    "requirement_node": "正在理解你的需求",
    "qa_node": "正在整理回答",
    "component_resource_node": "正在选择合适的图表和组件",
    "planner_node": "正在设计布局",
    "action_node": "正在准备画布修改",
    "validator_node": "正在检查修改结果",
    "reflection_node": "正在调整修改方案",
}


@contextmanager
def node_progress(node: str, state):
    """每次调用独立计时和编号，同步与异步入口共用"""

    try:
        # 读取当前图执行配置，用于判断是否处于图上下文，并获取任务标识
        config = get_config()
    except RuntimeError:
        # 独立调用节点没有图上下文，直接执行业务逻辑
        # 将控制权交给 with 块中的节点逻辑，此分支不发送进度事件
        yield
        # 节点逻辑正常执行后直接退出，不再获取图的流写入器
        return

    # 获取当前图的 custom 流写入器，供本次节点执行推送进度事件
    writer = get_stream_writer()
    # 以实际检查点配置为准，兼容没有在 state 中填写 ID 的本地示例
    thread_id = config.get("configurable", {}).get("thread_id") or state.get(
        "thread_id"
    )
    # 为本次节点调用生成唯一编号，区分同一节点的首次执行、重试和恢复执行
    execution_id = str(uuid4())
    # 使用单调计时器记录起点，后续计算本次执行耗时
    started = perf_counter()

    def emit(status: NodeProgressStatus):
        # 将任务、用户消息和本次执行编号放入同一事件，便于消费方关联进度
        event = NodeProgress(
            thread_id=thread_id,
            message_id=state.get("message_id"),
            execution_id=execution_id,
            node=node,
            # 未配置用户说明时返回 None，序列化时省略该字段
            label=NODE_PROGRESS_MESSAGES.get(node),
            status=status,
            # 开始事件不携带耗时，其他状态携带从本次调用开始计算的毫秒数
            duration_ms=None
            if status == "started"
            else round((perf_counter() - started) * 1000, 3),
        )

        # 转成可序列化的字典并发送到 custom 流，省略空字段，不写入业务状态
        writer(event.model_dump(mode="json", exclude_none=True))

    # 节点业务逻辑开始前，先通知消费方本次执行已经启动
    emit("started")
    try:
        # 执行 with 块中的节点逻辑，节点抛出的异常会回到下面对应的异常分支
        yield
    except GraphInterrupt:
        # interrupt 是正常暂停，恢复后会产生新的执行 ID
        emit("waiting_for_input")
        # 将暂停信号重新抛给 LangGraph，由图负责保存检查点和中断执行
        raise
    except asyncio.CancelledError:
        # 协程被取消也要结束这次节点执行，继续向上传递取消信号
        emit("cancelled")
        raise
    except Exception:
        # 只通知失败，异常详情仍由原有日志和 API 异常处理负责
        emit("failed")
        raise
    else:
        # 校验不通过等业务结果仍属于节点正常执行完成
        # completed 仅表示本次节点调用未抛异常，不代表整个任务成功或前端已执行
        emit("completed")