from langgraph.graph import StateGraph, START, END
from langgraph.graph.graph import CompiledGraph
from states import State
from tools.database_tools import db_query_tool
from tools.tool_nodes import query_gen_node, should_continue, final_answer_node
from tools.error_handling import create_tool_node_with_fallback

def run() -> CompiledGraph:
    workflow = StateGraph(State)
    workflow.add_node("query_gen", query_gen_node)
    workflow.add_node("execute_query", create_tool_node_with_fallback([db_query_tool]))
    workflow.add_node('final_answer', final_answer_node)

    workflow.add_edge(START, "query_gen")
    workflow.add_edge("query_gen", "execute_query")
    workflow.add_conditional_edges("execute_query", should_continue)
    workflow.add_edge("final_answer", END)

    # Compile the workflow into a runnable
    app = workflow.compile()

    return app