from IPython.display import Image, display
from langgraph.graph import StateGraph, END, START
from langgraph.graph.state import CompiledStateGraph
from graph import State, Nodes, Edges


class GraphBuilder:
    def __init__(self):
        """
        Initializes the GraphBuilder.
        """
        self.nodes = Nodes()
        self.edges = Edges()
        self.compiled_graph = None
        
    def build_agent_graph(self) -> CompiledStateGraph:
        """Build and return the agent graph."""
        graph = StateGraph(State)
        
        graph.add_node("manager", self.nodes.manager_node)
        graph.add_node("final_answer", self.nodes.final_answer_node)
        graph.add_node("auditor", self.nodes.auditor_node)
        graph.add_node("solver", self.nodes.solver_node)
        graph.add_node("researcher", self.nodes.researcher_node)
        graph.add_node("reasoner", self.nodes.reasoner_node)
        graph.add_node("viewer", self.nodes.viewer_node)
        
        graph.add_edge(START, "manager")
        graph.add_edge("final_answer", END)
        graph.add_edge("auditor", "manager")
        graph.add_edge("researcher", "solver")
        graph.add_edge("reasoner", "solver")
        graph.add_edge("viewer", "solver")

        graph.add_conditional_edges(
            "manager",
            self.edges.manager_edge,
            {
                "solver": "solver", "auditor": "auditor", "final_answer": "final_answer"
            }
        )
        graph.add_conditional_edges(
            "solver",
            self.edges.solver_edge,
            {
                "manager": "manager", "researcher": "researcher", "reasoner": "reasoner", "viewer": "viewer"
            }
        )
        
        return graph.compile()

    def display_graph(self):
        if self.compiled_graph is None:
            raise ValueError("Graph has not been compiled yet. Call build_agent_graph() first.")

        display(Image(self.compiled_graph.get_graph().draw_mermaid_png()))