from __future__ import annotations

from langgraph.graph import END, StateGraph

from agents import (
    claims_exception_node,
    claims_extraction_node,
    claims_intake_node,
    claims_normalization_node,
    claims_policy_similarity_node,
    claims_validation_node,
    claims_validation_router,
    scheduling_availability_node,
    scheduling_final_node,
    scheduling_parse_node,
    scheduling_provider_match_node,
    scheduling_readiness_node,
    scheduling_readiness_router,
)
from schemas import ClaimsState, SchedulingState


def build_claims_graph():
    graph = StateGraph(ClaimsState)
    graph.add_node("intake", claims_intake_node)
    graph.add_node("extraction", claims_extraction_node)
    graph.add_node("normalization", claims_normalization_node)
    graph.add_node("validation", claims_validation_node)
    graph.add_node("policy_similarity", claims_policy_similarity_node)
    graph.add_node("exception", claims_exception_node)

    graph.set_entry_point("intake")
    graph.add_edge("intake", "extraction")
    graph.add_edge("extraction", "normalization")
    graph.add_edge("normalization", "validation")
    graph.add_conditional_edges(
        "validation",
        claims_validation_router,
        {
            "retry_extraction": "extraction",
            "policy_similarity": "policy_similarity",
            "exception": "exception",
        },
    )
    graph.add_edge("policy_similarity", "exception")
    graph.add_edge("exception", END)
    return graph.compile()


def build_scheduling_graph():
    graph = StateGraph(SchedulingState)
    graph.add_node("parse", scheduling_parse_node)
    graph.add_node("readiness", scheduling_readiness_node)
    graph.add_node("provider_match", scheduling_provider_match_node)
    graph.add_node("availability", scheduling_availability_node)
    graph.add_node("final", scheduling_final_node)

    graph.set_entry_point("parse")
    graph.add_edge("parse", "readiness")
    graph.add_conditional_edges(
        "readiness",
        scheduling_readiness_router,
        {
            "provider_match": "provider_match",
            "final": "final",
        },
    )
    graph.add_edge("provider_match", "availability")
    graph.add_edge("availability", "final")
    graph.add_edge("final", END)
    return graph.compile()


claims_graph = build_claims_graph()
scheduling_graph = build_scheduling_graph()