from langgraph.graph import StateGraph
from builder.state import AgentState
from builder.nodes import sql_generator, sql_executor, chart_generator, narrator, decider, general_chat

def build_graph():
    graph = StateGraph(AgentState)

    graph.add_node("decider", decider)
    graph.add_node("sql_generator", sql_generator)
    graph.add_node("sql_executor", sql_executor)
    graph.add_node("chart_generator", chart_generator)
    graph.add_node("narrator", narrator)
    graph.add_node("general_chat", general_chat)

    graph.set_entry_point("decider")

    graph.add_conditional_edges(
        "decider",
        lambda state: state["route"],
        {
            "sql": "sql_generator",
            "chat": "general_chat",
        },
    )

    graph.add_edge("sql_generator", "sql_executor")
    graph.add_edge("sql_executor", "chart_generator")
    graph.add_edge("sql_executor", "narrator")

    return graph.compile()