from __future__ import annotations

from typing import Any
from pydantic import ConfigDict
from langchain_core.documents import Document
from langchain_core.retrievers import BaseRetriever

from .core import BayesianIndex, RetrievalHit


class BayesianRetriever(BaseRetriever):
    """LangChain-native retriever backed by the generalized Bayesian index.

    The LangChain contract returns ``Document`` objects.  All posterior fields
    required downstream are copied into document metadata so an end-to-end RAG
    chain can use the *actual* retriever output without recomputing retrieval.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)
    index: Any
    k: int = 4

    def search_hits(self, query: str) -> list[RetrievalHit]:
        return self.index.search(query, k=self.k)

    @staticmethod
    def _hit_to_document(h: RetrievalHit) -> Document:
        md = dict(h.document.metadata)
        md.update({
            "id": h.document.id,
            "title": h.document.title,
            "source": h.document.source,
            "retrieval_score": h.retrieval_score,
            "log_prior": h.log_prior,
            "log_likelihood": h.log_likelihood,
            "log_posterior_unnormalized": h.log_posterior_unnormalized,
            "posterior": h.posterior,
            "query_coverage": h.query_coverage,
            "rank": h.rank,
        })
        return Document(page_content=h.document.text, metadata=md)

    def _get_relevant_documents(self, query: str, *, run_manager: Any = None) -> list[Document]:
        return [self._hit_to_document(h) for h in self.search_hits(query)]