"""
graph.py

Builds the core WalkXR LangGraph workflow for a single turn.
"""

from __future__ import annotations
from langgraph.graph import END, START, StateGraph
from walkxr_ai.core.state import WalkState
from walkxr_ai.core.routing import (
    route_after_call_llm,
    route_after_supervisor_score,
    route_after_update_conversation_history,
)
from walkxr_ai.core.nodes import (
    build_call_llm_node,
    build_retrieve_context_node,
    construct_prompt,
    update_conversation_history,
    build_call_supervisor_node,
    build_update_memory_node,
    fail_safe_response,
)


def build_graph(retrieval_engine, user_memory_store, responder, memory_extractor):
    """
    Build and compile WalkXR LangGraph for one conversation turn.
    """
    graph = StateGraph(WalkState)

    retrieve_context = build_retrieve_context_node(retrieval_engine, user_memory_store)
    call_llm = build_call_llm_node(responder)
    call_supervisor = build_call_supervisor_node(responder)
    update_memory = build_update_memory_node(memory_extractor)

    graph.add_node("retrieve_context", retrieve_context)
    graph.add_node("construct_prompt", construct_prompt)
    graph.add_node("call_llm", call_llm)
    graph.add_node("call_supervisor", call_supervisor)
    graph.add_node("fail_safe_response", fail_safe_response)
    graph.add_node("update_conversation_history", update_conversation_history)
    graph.add_node("update_memory", update_memory)

    graph.add_edge(START, "retrieve_context")
    graph.add_edge("retrieve_context", "construct_prompt")
    graph.add_edge("construct_prompt", "call_llm")
    graph.add_conditional_edges("call_llm", route_after_call_llm,
        {
            "call_supervisor": "call_supervisor",
            "error": END,
        },
    )
    graph.add_conditional_edges("call_supervisor", route_after_supervisor_score,
        {
            "loop": "call_llm",
            "fail-safe": "fail_safe_response",
            "exit": "update_conversation_history"
        }
    )
    graph.add_edge("fail_safe_response", "update_conversation_history")
    graph.add_conditional_edges("update_conversation_history", route_after_update_conversation_history,
        {
            "update_memory": "update_memory",
            "end": END,
        },
    )
    graph.add_edge("update_memory", END)

    return graph.compile()