from langgraph.graph import StateGraph, END
from .nodes import (
    ChatBotState,
    classifier_agent,
    general_info_agent,
    catalog_agent,
    order_agent,
)


def create_chat_graph():
    # Crear el grafo con el estado tipado
    graph = StateGraph(ChatBotState)

    # Añadir nodos
    graph.add_node("classifier", classifier_agent)
    graph.add_node("general_info", general_info_agent)
    graph.add_node("catalog", catalog_agent)
    graph.add_node("order", order_agent)

    # Configurar el punto de entrada
    graph.set_entry_point("classifier")

    # Función de enrutamiento basada en la categoría
    def route_to_agent(state):
        return state["category"]

    # Añadir bordes condicionales
    graph.add_conditional_edges(
        "classifier",
        route_to_agent,
        {
            "general_info": "general_info",
            "catalog": "catalog",
            "order": "order",
        },
    )

    # Conectar al nodo final
    graph.add_edge("general_info", END)
    graph.add_edge("catalog", END)
    graph.add_edge("order", END)

    return graph.compile()