"""输入 → 意图（是不是明确的问答）→ 问答节点或 screen agent；agent 追问时暂停等待用户"""

import inspect

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
from langgraph.graph import END, StateGraph

from app.processor.nodes.input_node import InputNode
from app.processor.nodes.intent_node import IntentNode
from app.processor.nodes.qa_node import QANode
from app.processor.nodes.screen_agent_node import ScreenAgentNode
from app.processor.nodes.wait_for_user_node import WaitForUserNode
from app.processor.state import QueryState, RunContext


def build_checkpoint_serde() -> JsonPlusSerializer:
    # 允许恢复状态中使用的业务模型类型，使检查点读取后仍可按模型对象访问字段
    # 内存与持久化检查点共用，否则持久化读回的状态无法还原为模型对象
    return JsonPlusSerializer(
        allowed_msgpack_modules=[
            ("app.schemas.input_context", "InputContext"),
            ("app.schemas.decision", "DecisionResponse"),
            ("app.schemas.actions.add", "AddAction"),
            ("app.schemas.actions.delete", "DeleteAction"),
            ("app.schemas.actions.move", "MoveAction"),
            ("app.schemas.actions.resize", "ResizeAction"),
            ("app.schemas.actions.update", "UpdateAction"),
            ("app.schemas.actions.update_canvas", "UpdateCanvasAction"),
            ("app.schemas.actions.set_interaction", "SetInteractionAction"),
            ("app.schemas.clarification", "Clarification"),
            ("app.schemas.clarification_reply", "ClarificationReply"),
        ]
    )


def route_after_intent(state: QueryState) -> str:
    return "qa_node" if state["intent"] == "qa" else "screen_agent_node"


def route_after_agent(state: QueryState) -> str:
    return "wait_for_user_node" if state["status"] == "waiting_for_input" else END


def route_after_wait(state: QueryState) -> str:
    # 用户回答后回到输入节点合并上下文，取消则结束
    return "input_node" if state["status"] == "running" else END


def build_graph(checkpointer=None):
    workflow = StateGraph(QueryState, context_schema=RunContext)
    nodes = {
        "input_node": InputNode(),
        "intent_node": IntentNode(),
        "qa_node": QANode(),
        "screen_agent_node": ScreenAgentNode(),
        "wait_for_user_node": WaitForUserNode(),
    }
    for node_name, node in nodes.items():
        # 异步节点使用 BaseNode 的异步入口，保留统一的日志和进度封装
        action = node._acall if inspect.iscoroutinefunction(node.process) else node
        workflow.add_node(node_name, action)

    workflow.set_entry_point("input_node")
    workflow.add_edge("input_node", "intent_node")
    workflow.add_conditional_edges(
        "intent_node",
        route_after_intent,
        {"qa_node": "qa_node", "screen_agent_node": "screen_agent_node"},
    )
    workflow.add_edge("qa_node", END)
    workflow.add_conditional_edges(
        "screen_agent_node",
        route_after_agent,
        {"wait_for_user_node": "wait_for_user_node", END: END},
    )
    workflow.add_conditional_edges(
        "wait_for_user_node", route_after_wait, {"input_node": "input_node", END: END}
    )

    # 调用方未提供检查点存储时使用进程内存，支持暂停恢复，进程退出后不保留
    if checkpointer is None:
        checkpointer = InMemorySaver(serde=build_checkpoint_serde())
    return workflow.compile(checkpointer=checkpointer)