from langgraph.graph import StateGraph, END
from .state import DevAgentState
from .nodes.analyze import analyze_requirements
from .nodes.clarify import clarify
from .nodes.generate import generate_spec
from .nodes.validate import validate_spec
from .nodes.finalize import finalize

def build_graph():
    graph = StateGraph(DevAgentState)
    graph.add_node('analyze', analyze_requirements)
    graph.add_node('clarify', clarify)
    graph.add_node('generate', generate_spec)
    graph.add_node('validate', validate_spec)
    graph.add_node('finalize', finalize)

    graph.set_entry_point('analyze')
    
    def route_analysis(state: DevAgentState):
        if state.get("status") == "awaiting_clarification":
            return "clarify"
        return "generate"

    graph.add_conditional_edges(
        'analyze',
        route_analysis,
        {
            'clarify': 'clarify',
            'generate': 'generate'
        }
    )
    
    graph.add_edge('clarify', END)
    graph.add_edge('generate', 'validate')
    
    def route_validation(state: DevAgentState):
        vr = state.get('validation_result', {}) or {}
        if vr.get('valid'):
            return 'finalize'
        else:
            return 'clarify'
            
    graph.add_conditional_edges(
        'validate',
        route_validation,
        {
            'finalize': 'finalize',
            'clarify': 'clarify'
        }
    )
    
    graph.add_edge('finalize', END)
    return graph.compile()