from langgraph.graph import StateGraph, START, END
try:
    from langgraph.checkpoint.memory import MemorySaver
except ImportError:
    from langgraph.checkpoint import MemorySaver

from src.agent.state import AgentState
from src.agent.nodes import (
    orchestrator_node,
    should_continue,
    masker_node,
    response_node,
    tool_node,
)


def build_graph(checkpointer=None) -> "CompiledGraph":
    g = StateGraph(AgentState)

    g.add_node("orchestrator", orchestrator_node)
    g.add_node("tools",        tool_node)
    g.add_node("masker",       masker_node)
    g.add_node("response",     response_node)

    g.add_edge(START, "orchestrator")

    g.add_conditional_edges(
        "orchestrator",
        should_continue,
        {"tools": "tools", "masker": "masker"},
    )

    g.add_edge("tools",    "orchestrator")  # loop back for more tool calls
    g.add_edge("masker",   "response")
    g.add_edge("response", END)

    return g.compile(checkpointer=checkpointer or MemorySaver())