from __future__ import annotations

from typing import Any, Sequence
from langchain_core.documents import Document
from langchain_core.runnables import RunnableLambda

from .core import BayesianIndex, DocumentRecord, ExtractiveGroundedGenerator, RetrievalHit
from .retriever import BayesianRetriever


class BayesianRAG:
    """End-to-end LangChain RAG with explicit posterior uncertainty.

    Retrieval is performed exactly once through ``BaseRetriever.invoke``.  The
    returned LangChain ``Document`` objects carry the posterior metadata and are
    converted back to typed hits only for the uncertainty/grounding layer.
    """

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

    @staticmethod
    def _documents_to_hits(documents: Sequence[Document]) -> list[RetrievalHit]:
        hits: list[RetrievalHit] = []
        for fallback_rank, doc in enumerate(documents, start=1):
            md = dict(doc.metadata)
            rec = DocumentRecord(
                id=str(md.pop("id", f"doc-{fallback_rank}")),
                text=doc.page_content,
                title=str(md.pop("title", "")),
                source=str(md.pop("source", "unknown")),
                metadata={
                    k: v
                    for k, v in md.items()
                    if k not in {
                        "retrieval_score", "log_prior", "log_likelihood",
                        "log_posterior_unnormalized", "posterior",
                        "query_coverage", "rank",
                    }
                },
            )
            hits.append(RetrievalHit(
                document=rec,
                retrieval_score=float(doc.metadata.get("retrieval_score", 0.0)),
                log_prior=float(doc.metadata.get("log_prior", 0.0)),
                log_likelihood=float(doc.metadata.get("log_likelihood", 0.0)),
                log_posterior_unnormalized=float(doc.metadata.get("log_posterior_unnormalized", 0.0)),
                posterior=float(doc.metadata.get("posterior", 0.0)),
                query_coverage=float(doc.metadata.get("query_coverage", 0.0)),
                rank=int(doc.metadata.get("rank", fallback_rank)),
            ))
        return hits

    def invoke(self, question: str) -> dict[str, Any]:
        documents = self.retriever.invoke(question)
        hits = self._documents_to_hits(documents)
        if self.index.should_abstain(hits):
            answer = ExtractiveGroundedGenerator.abstain_text
            citations: list[str] = []
        elif hasattr(self.generator, "generate"):
            answer, citations = self.generator.generate(question, 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: {question}"
            )
            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 {
            "question": question,
            "answer": answer,
            "citations": citations,
            "abstained": not citations,
            "hits": [h.as_dict() for h in hits],
        }

    def as_runnable(self) -> RunnableLambda:
        return RunnableLambda(self.invoke)