from langgraph.graph import StateGraph, END
from ea_chatbot.graph.state import AgentState
from ea_chatbot.graph.nodes.query_analyzer import query_analyzer_node
from ea_chatbot.graph.nodes.planner import planner_node
from ea_chatbot.graph.nodes.coder import coder_node
from ea_chatbot.graph.nodes.error_corrector import error_corrector_node
from ea_chatbot.graph.nodes.executor import executor_node
from ea_chatbot.graph.nodes.summarizer import summarizer_node
from ea_chatbot.graph.nodes.researcher import researcher_node
from ea_chatbot.graph.nodes.clarification import clarification_node
from ea_chatbot.graph.nodes.summarize_conversation import summarize_conversation_node

MAX_ITERATIONS = 3

def router(state: AgentState) -> str:
    """Route to the next node based on the analysis."""
    next_action = state.get("next_action")
    if next_action == "plan":
        return "planner"
    elif next_action == "research":
        return "researcher"
    elif next_action == "clarify":
        return "clarification"
    else:
        return END

def create_workflow():
    """Create the LangGraph workflow."""
    workflow = StateGraph(AgentState)
    
    # Add nodes
    workflow.add_node("query_analyzer", query_analyzer_node)
    workflow.add_node("planner", planner_node)
    workflow.add_node("coder", coder_node)
    workflow.add_node("error_corrector", error_corrector_node)
    workflow.add_node("researcher", researcher_node)
    workflow.add_node("clarification", clarification_node)
    workflow.add_node("executor", executor_node)
    workflow.add_node("summarizer", summarizer_node)
    workflow.add_node("summarize_conversation", summarize_conversation_node)
    
    # Set entry point
    workflow.set_entry_point("query_analyzer")
    
    # Add conditional edges from query_analyzer
    workflow.add_conditional_edges(
        "query_analyzer",
        router,
        {
            "planner": "planner",
            "researcher": "researcher",
            "clarification": "clarification",
            END: END
        }
    )
    
    # Linear flow for planning and coding
    workflow.add_edge("planner", "coder")
    workflow.add_edge("coder", "executor")
    
    # Executor routing
    def executor_router(state: AgentState) -> str:
        if state.get("error"):
            # Check for iteration limit to prevent infinite loops
            if state.get("iterations", 0) >= MAX_ITERATIONS:
                return "summarizer"
            return "error_corrector"
        return "summarizer"
        
    workflow.add_conditional_edges(
        "executor",
        executor_router,
        {
            "error_corrector": "error_corrector",
            "summarizer": "summarizer"
        }
    )
    
    workflow.add_edge("error_corrector", "executor")
    
    workflow.add_edge("researcher", "summarize_conversation")
    workflow.add_edge("clarification", END)
    workflow.add_edge("summarizer", "summarize_conversation")
    workflow.add_edge("summarize_conversation", END)
    
    # Compile the graph
    app = workflow.compile()
    
    return app

# Initialize the app
app = create_workflow()