"""主图与 Attempt 子图只负责节点、边和条件路由注册。"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any, Literal, cast

from langgraph.graph import END, START, StateGraph
from langgraph.types import Send

from app.graph.attempt_graph_node.evaluate_judge import make_evaluate_judge
from app.graph.attempt_graph_node.evaluate_rules import make_evaluate_rules
from app.graph.attempt_graph_node.invoke_provider import make_invoke_provider
from app.graph.attempt_graph_node.parse_output import make_parse_output
from app.graph.attempt_graph_node.persist_attempt import make_persist_attempt
from app.graph.attempt_graph_node.prepare_attempt import make_prepare_attempt
from app.graph.attempt_graph_node.validate_schema import make_validate_schema
from app.graph.graph_node.aggregate import make_aggregate
from app.graph.graph_node.apply_gates import make_apply_gates
from app.graph.graph_node.begin_scored import make_begin_scored
from app.graph.graph_node.compute_pareto import make_compute_pareto
from app.graph.graph_node.human_review import make_human_review
from app.graph.graph_node.load_snapshot import make_load_snapshot
from app.graph.graph_node.plan_attempts import make_plan_attempts
from app.graph.graph_node.preflight import make_preflight
from app.graph.graph_node.publish_report import make_publish_report
from app.graph.state import AttemptGraphState, RunGraphState
from app.services.workflow import WorkflowService

_ATTEMPT_TERMINAL = {"succeeded", "failed", "uncertain", "cancelled"}


def _node(action: Any) -> Any:
    """隔离 LangGraph 第三方类型重载差异，运行时仍保留完整节点签名。"""

    return action


def _continue_or_persist(
    next_node: str,
) -> Callable[[AttemptGraphState], str]:
    def route(state: AttemptGraphState) -> str:
        return "persist_attempt" if state.get("attempt_status") in _ATTEMPT_TERMINAL else next_node

    return route


def build_attempt_graph(service: WorkflowService) -> Any:
    graph = StateGraph(AttemptGraphState)
    graph.add_node("prepare_attempt", _node(make_prepare_attempt(service)))
    graph.add_node("invoke_provider", _node(make_invoke_provider(service)))
    graph.add_node("parse_output", _node(make_parse_output(service)))
    graph.add_node("validate_schema", _node(make_validate_schema(service)))
    graph.add_node("evaluate_rules", _node(make_evaluate_rules(service)))
    graph.add_node("evaluate_judge", _node(make_evaluate_judge(service)))
    graph.add_node("persist_attempt", _node(make_persist_attempt(service)))
    graph.add_edge(START, "prepare_attempt")
    graph.add_conditional_edges(
        "prepare_attempt", _continue_or_persist("invoke_provider")
    )
    graph.add_conditional_edges(
        "invoke_provider", _continue_or_persist("parse_output")
    )
    graph.add_conditional_edges(
        "parse_output", _continue_or_persist("validate_schema")
    )
    graph.add_conditional_edges(
        "validate_schema", _continue_or_persist("evaluate_rules")
    )
    graph.add_conditional_edges(
        "evaluate_rules", _continue_or_persist("evaluate_judge")
    )
    graph.add_conditional_edges(
        "evaluate_judge", _continue_or_persist("persist_attempt")
    )
    graph.add_edge("persist_attempt", END)
    return graph.compile()


def _after_load(state: RunGraphState) -> Literal["preflight", "__end__"]:
    return cast(
        Literal["preflight", "__end__"],
        END if state.get("status") == "cancelled" else "preflight",
    )


def _send_attempts(node: str, state: RunGraphState, attempts: list[str]) -> list[Send]:
    return [
        Send(
            node,
            {
                "run_id": state["run_id"],
                "attempt_id": attempt_id,
                "attempt_status": "planned",
                "completed_attempt_ids": [],
            },
        )
        for attempt_id in attempts
    ]


def _dispatch_warmups(
    state: RunGraphState,
) -> list[Send] | Literal["begin_scored"]:
    attempts = state.get("warmup_attempt_ids", [])
    return _send_attempts("run_warmup", state, attempts) if attempts else "begin_scored"


def _dispatch_attempts(state: RunGraphState) -> list[Send] | Literal["aggregate"]:
    attempts = state.get("attempt_ids", [])
    if not attempts:
        return "aggregate"
    return _send_attempts("run_attempt", state, attempts)


def _after_aggregate(state: RunGraphState) -> Literal["human_review", "apply_gates"]:
    return "human_review" if state.get("human_review", False) else "apply_gates"


def build_run_graph(service: WorkflowService, checkpointer: Any | None = None) -> Any:
    warmup_graph = build_attempt_graph(service)
    attempt_graph = build_attempt_graph(service)
    graph = StateGraph(RunGraphState)
    graph.add_node("load_snapshot", _node(make_load_snapshot(service)))
    graph.add_node("preflight", _node(make_preflight(service)))
    graph.add_node("plan_attempts", _node(make_plan_attempts(service)))
    graph.add_node("run_warmup", _node(warmup_graph))
    graph.add_node("begin_scored", _node(make_begin_scored()))
    graph.add_node("run_attempt", _node(attempt_graph))
    graph.add_node("aggregate", _node(make_aggregate(service)))
    graph.add_node("human_review", _node(make_human_review(service)))
    graph.add_node("apply_gates", _node(make_apply_gates(service)))
    graph.add_node("compute_pareto", _node(make_compute_pareto(service)))
    graph.add_node("publish_report", _node(make_publish_report(service)))

    graph.add_edge(START, "load_snapshot")
    graph.add_conditional_edges("load_snapshot", _after_load)
    graph.add_edge("preflight", "plan_attempts")
    graph.add_conditional_edges("plan_attempts", _dispatch_warmups)
    graph.add_edge("run_warmup", "begin_scored")
    graph.add_conditional_edges("begin_scored", _dispatch_attempts)
    graph.add_edge("run_attempt", "aggregate")
    graph.add_conditional_edges("aggregate", _after_aggregate)
    graph.add_edge("human_review", "apply_gates")
    graph.add_edge("apply_gates", "compute_pareto")
    graph.add_edge("compute_pareto", "publish_report")
    graph.add_edge("publish_report", END)
    return graph.compile(checkpointer=checkpointer)


__all__ = ["build_attempt_graph", "build_run_graph"]