"""
graph.py
---------
Wires the five agents into a real LangGraph StateGraph implementing:

    Planner -> Retriever -> Evidence -> Reasoning -> Verification
                   ^                                      |
                   |__________ (insufficient, iterate) ___|
                                                            |
                                                  (sufficient / max iters)
                                                            v
                                                      Final Answer

This is the "PLAN -> RETRIEVE -> CHECK -> REASON -> VERIFY -> RETRIEVE AGAIN
IF NECESSARY -> ANSWER" loop from the spec, expressed as an actual
conditional edge in LangGraph rather than a fixed linear pipeline.
"""
from __future__ import annotations

from langgraph.graph import END, StateGraph

from src.agents.evidence import evaluate_evidence
from src.agents.final_answer import finalize
from src.agents.planner import plan
from src.agents.reasoning import reason
from src.agents.retrieval import retrieve
from src.agents.state import ResearchState
from src.agents.verification import verify


def _route_after_verification(state: ResearchState) -> str:
    if state.get("error"):
        return "end"
    if state.get("verification_passed"):
        return "finalize"
    # Not passed but verify() only leaves verification_passed=False when a retry is queued
    return "retrieve"


def _route_after_plan(state: ResearchState) -> str:
    return "end" if state.get("error") else "retrieve"


def build_graph():
    graph = StateGraph(ResearchState)

    graph.add_node("plan", plan)
    graph.add_node("retrieve", retrieve)
    graph.add_node("evaluate_evidence", evaluate_evidence)
    graph.add_node("reason", reason)
    graph.add_node("verify", verify)
    graph.add_node("finalize", finalize)

    graph.set_entry_point("plan")
    graph.add_conditional_edges("plan", _route_after_plan, {"retrieve": "retrieve", "end": END})
    graph.add_edge("retrieve", "evaluate_evidence")
    graph.add_edge("evaluate_evidence", "reason")
    graph.add_edge("reason", "verify")
    graph.add_conditional_edges(
        "verify", _route_after_verification, {"retrieve": "retrieve", "finalize": "finalize", "end": END}
    )
    graph.add_edge("finalize", END)

    return graph.compile()


_COMPILED_GRAPH = None


def get_graph():
    global _COMPILED_GRAPH
    if _COMPILED_GRAPH is None:
        _COMPILED_GRAPH = build_graph()
    return _COMPILED_GRAPH


def run_research(question: str, max_iterations: int = 3) -> ResearchState:
    initial_state: ResearchState = {
        "question": question,
        "max_iterations": max_iterations,
        "iteration": 0,
        "trace": [],
    }
    graph = get_graph()
    # recursion_limit accounts for the loop: each iteration touches ~4 nodes
    result = graph.invoke(initial_state, config={"recursion_limit": 6 * max_iterations + 10})
    return result