from langgraph.graph import StateGraph, START, END
from .utils.state import State
from .utils.nodes import ideator , critic , improver , validator1 , validator2 , validator3 , validator4 , route1_after_validation , route2_after_validation , route3_after_validation , route4_after_validation


def ideation_graph():
    graph_builder= StateGraph(State)
    graph_builder.add_node(ideator)
    graph_builder.add_node(critic)
    graph_builder.add_node(improver)
    graph_builder.add_node(validator1)
    graph_builder.add_node(validator2)
    graph_builder.add_node(validator3)
    graph_builder.add_node(validator4)


    graph_builder.add_edge(START, "ideator")  # Start the graph with node_1
    graph_builder.add_edge("ideator", "critic")
    graph_builder.add_edge("critic", "improver")
    graph_builder.add_edge("improver", "validator1")
    graph_builder.add_edge("validator1", "validator2")
    graph_builder.add_edge("validator2", "validator3")
    graph_builder.add_edge("validator3", "validator4")
    graph_builder.add_edge("validator4", END)

    # Use conditional routing from validator
    graph_builder.add_conditional_edges("validator1", route1_after_validation,{False:'improver',True:'validator2'})
    graph_builder.add_conditional_edges("validator2", route2_after_validation,{False:'improver',True:'validator3'})
    graph_builder.add_conditional_edges("validator3", route3_after_validation,{False:'improver',True:'validator4'})
    graph_builder.add_conditional_edges("validator4", route4_after_validation,{False:'improver',True:END})

    return graph_builder.compile()