from langgraph.graph import StateGraph, START, END , MessagesState
from .utils.state import State
from .utils.nodes import RetrieverNode, IdeatorNode , ModeratorNode ,NormalizerNode, ValidatorNode , RoutingAfterValidation, JudgeNode1 , JudgeNode2, Aggregrator, SimplifierNode
from langgraph.checkpoint.memory import MemorySaver
import os


class IdeationAgent:
    def __init__(self):
        self.memory = MemorySaver()

    def ideation_graph(self):
        graph_builder= StateGraph(State)
        graph_builder.add_node("retriever", RetrieverNode().run)
        graph_builder.add_node("ideator", IdeatorNode().run)
        graph_builder.add_node("moderator", ModeratorNode().run)
        graph_builder.add_node("simplifier", SimplifierNode().run)
        graph_builder.add_node("normalizer",NormalizerNode().run)
        graph_builder.add_node("judge1", JudgeNode1().run)
        graph_builder.add_node("judge2", JudgeNode2().run)
        graph_builder.add_node("aggregrator", Aggregrator().run)
        graph_builder.add_node("validator", ValidatorNode().run)

        graph_builder.add_edge(START, "retriever")
        graph_builder.add_edge("retriever", "ideator")  
        graph_builder.add_edge("ideator", "moderator")  
        graph_builder.add_edge("moderator", "simplifier")
        graph_builder.add_edge("simplifier", "normalizer")
        graph_builder.add_edge("normalizer", "judge1")
        graph_builder.add_edge("normalizer","judge2")
        graph_builder.add_edge("judge1", "aggregrator")
        graph_builder.add_edge("judge2", "aggregrator")
        graph_builder.add_edge("aggregrator", END)
        


        # graph_builder.add_conditional_edges("validator", RoutingAfterValidation().route,{False:'critic',True:END})

        return graph_builder.compile(checkpointer=self.memory)