from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import InMemorySaver

from .state import AgentState
from .extractor import extract_parameters
from .selector import select_tool
from .validator import validate_tool
from .clarification import ask_user
from .executor import execute_tool
from .response import generate_response


def route_after_validation(state: AgentState):
    if state["selected_tool"] is None:
        return "no_tool"

    if state["missing_params"]:
        return "ask_user"

    return "execute_tool"


builder = StateGraph(AgentState)

# Register nodes
builder.add_node("extract", extract_parameters)
builder.add_node("select_tool", select_tool)
builder.add_node("validate", validate_tool)
builder.add_node("ask_user", ask_user)
builder.add_node("execute_tool", execute_tool)
builder.add_node("generate_response", generate_response)

# Main flow
builder.add_edge(START, "extract")
builder.add_edge("extract", "select_tool")
builder.add_edge("select_tool", "validate")

# Conditional routing
builder.add_conditional_edges(
    "validate",
    route_after_validation,
    {
        "no_tool": END,
        "ask_user": "ask_user",
        "execute_tool": "execute_tool"
    }
)

# Final edges
builder.add_edge("ask_user", END)
builder.add_edge("execute_tool", "generate_response")
builder.add_edge("generate_response", END)

graph = builder.compile(
    checkpointer=InMemorySaver()
)