"""LangGraph graph definition and compilation."""

from dotenv import load_dotenv
from langgraph.graph import END, StateGraph

from src.nodes.call_llm import call_llm
from src.nodes.construct_prompt import construct_prompt
from src.nodes.reflection import reflection
from src.nodes.retrieve_context import retrieve_context
from src.nodes.update_memory import update_memory
from src.routing import route_after_memory_update
from src.schemas.state import WalkState

load_dotenv()


def build_graph() -> StateGraph:
    """Build and compile the WalkXR role-play graph."""

    graph = StateGraph(WalkState)

    # Register nodes
    graph.add_node("retrieve_context", retrieve_context)
    graph.add_node("construct_prompt", construct_prompt)
    graph.add_node("call_llm", call_llm)
    graph.add_node("update_memory", update_memory)
    graph.add_node("reflection", reflection)

    # Connect edges
    graph.set_entry_point("retrieve_context")
    graph.add_edge("retrieve_context", "construct_prompt")
    graph.add_edge("construct_prompt", "call_llm")
    graph.add_edge("call_llm", "update_memory")

    # Conditional branching after update_memory
    graph.add_conditional_edges(
        "update_memory",
        route_after_memory_update,
        {
            "reflection": "reflection",
            "end": END,
        },
    )

    # reflection -> END
    graph.add_edge("reflection", END)

    return graph.compile()


# Singleton instance
app = build_graph()