from __future__ import annotations

from typing import Literal

from langgraph.graph import END, START, StateGraph

from .agents import orchestrator_node, requirements_analyst_node, reviewer_node, test_designer_node
from .state import QAState


def route_from_orchestrator(state: QAState) -> Literal["finish", "requirements_analyst", "test_designer", "reviewer", "fail"]:
    decision = state.get("orchestration_decision", {})
    if state.get("errors") and decision.get("decision") != "backtrack":
        return "fail"
    if decision.get("decision") == "finish":
        return "finish"
    if decision.get("decision") == "backtrack":
        target = decision.get("target_agent")
        if target in {"requirements_analyst", "test_designer", "reviewer"}:
            return target
    return "fail"


def build_graph():
    graph = StateGraph(QAState)
    graph.add_node("requirements_analyst", requirements_analyst_node)
    graph.add_node("test_designer", test_designer_node)
    graph.add_node("reviewer", reviewer_node)
    graph.add_node("orchestrator", orchestrator_node)
    graph.add_edge(START, "requirements_analyst")
    graph.add_edge("requirements_analyst", "test_designer")
    graph.add_edge("test_designer", "reviewer")
    graph.add_edge("reviewer", "orchestrator")
    graph.add_conditional_edges(
        "orchestrator",
        route_from_orchestrator,
        {
            "finish": END,
            "requirements_analyst": "requirements_analyst",
            "test_designer": "test_designer",
            "reviewer": "reviewer",
            "fail": END,
        },
    )
    return graph.compile()