from __future__ import annotations

from typing import Any, Literal, TypedDict

from .core import BayesianIndex, ExtractiveGroundedGenerator, RetrievalHit


class RAGState(TypedDict, total=False):
    query: str
    hits: list[RetrievalHit]
    decision: Literal["generate", "abstain"]
    answer: str
    citations: list[str]
    abstained: bool


class BayesianRAGGraph:
    """Actual LangGraph orchestration around the Bayesian retriever.

    Graph topology:
        START -> retrieve -> assess -> {generate | abstain} -> END
    """

    def __init__(self, index: BayesianIndex, generator: Any | None = None, k: int = 4):
        self.index = index
        self.generator = generator or ExtractiveGroundedGenerator()
        self.k = k
        self.graph = self._build_graph()

    def _retrieve(self, state: RAGState) -> dict[str, Any]:
        return {"hits": self.index.search(state["query"], k=self.k)}

    def _assess(self, state: RAGState) -> dict[str, Any]:
        decision: Literal["generate", "abstain"] = "abstain" if self.index.should_abstain(state.get("hits", [])) else "generate"
        return {"decision": decision}

    @staticmethod
    def _route(state: RAGState) -> str:
        return state.get("decision", "abstain")

    def _generate(self, state: RAGState) -> dict[str, Any]:
        hits = state.get("hits", [])
        if hasattr(self.generator, "generate"):
            answer, citations = self.generator.generate(state["query"], hits)
        else:
            context = "\n".join(f"[{h.document.id}] {h.document.text}" for h in hits)
            prompt = (
                "Answer only from the supplied context. Cite document ids in square brackets. "
                "If unsupported, say you do not know.\n\n"
                f"Context:\n{context}\n\nQuestion: {state['query']}"
            )
            resp = self.generator.invoke(prompt)
            answer = getattr(resp, "content", str(resp))
            citations = [h.document.id for h in hits if f"[{h.document.id}]" in answer]
        return {"answer": answer, "citations": citations, "abstained": False}

    @staticmethod
    def _abstain(state: RAGState) -> dict[str, Any]:
        return {
            "answer": ExtractiveGroundedGenerator.abstain_text,
            "citations": [],
            "abstained": True,
        }

    def _build_graph(self):
        try:
            from langgraph.graph import StateGraph, START, END
        except ImportError as exc:
            raise ImportError("bayesian-langgraph-rag requires langgraph>=1.2,<2.0") from exc

        builder = StateGraph(RAGState)
        builder.add_node("retrieve", self._retrieve)
        builder.add_node("assess", self._assess)
        builder.add_node("generate", self._generate)
        builder.add_node("abstain", self._abstain)
        builder.add_edge(START, "retrieve")
        builder.add_edge("retrieve", "assess")
        builder.add_conditional_edges(
            "assess",
            self._route,
            {"generate": "generate", "abstain": "abstain"},
        )
        builder.add_edge("generate", END)
        builder.add_edge("abstain", END)
        return builder.compile()

    def invoke(self, query: str) -> dict[str, Any]:
        state = self.graph.invoke({"query": query})
        hits = state.get("hits", [])
        return {
            "question": query,
            "answer": state.get("answer", ExtractiveGroundedGenerator.abstain_text),
            "citations": state.get("citations", []),
            "abstained": state.get("abstained", True),
            "hits": [h.as_dict() for h in hits],
        }