from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Iterable, Mapping, Sequence
import math
import re

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer, ENGLISH_STOP_WORDS


def _clip01(x: float, eps: float = 1e-6) -> float:
    return min(max(float(x), eps), 1.0)


def _logsumexp(xs: Sequence[float]) -> float:
    if not xs:
        return float('-inf')
    finite = [x for x in xs if math.isfinite(x)]
    if not finite:
        return float('-inf')
    m = max(finite)
    return m + math.log(sum(math.exp(x - m) for x in finite))


@dataclass(frozen=True)
class DocumentRecord:
    id: str
    text: str
    title: str = ""
    source: str = "unknown"
    metadata: Mapping[str, Any] = field(default_factory=dict)

    @property
    def retrieval_text(self) -> str:
        return (self.title + ". " + self.text).strip(". ")

    def source_prior(self) -> float:
        md = self.metadata
        a = max(float(md.get("reliability_alpha", 1.0)), 1e-6)
        b = max(float(md.get("reliability_beta", 1.0)), 1e-6)
        reliability = a / (a + b)
        authority = _clip01(md.get("authority", 0.5))
        freshness = _clip01(md.get("freshness", 0.5))
        return reliability * authority * freshness


@dataclass(frozen=True)
class RetrievalHit:
    document: DocumentRecord
    retrieval_score: float
    log_prior: float
    log_likelihood: float
    log_posterior_unnormalized: float
    posterior: float
    query_coverage: float
    rank: int

    def as_dict(self) -> dict[str, Any]:
        return {
            "id": self.document.id,
            "title": self.document.title,
            "text": self.document.text,
            "source": self.document.source,
            "metadata": dict(self.document.metadata),
            "retrieval_score": self.retrieval_score,
            "log_prior": self.log_prior,
            "log_likelihood": self.log_likelihood,
            "posterior": self.posterior,
            "query_coverage": self.query_coverage,
            "rank": self.rank,
        }


@dataclass(frozen=True)
class BayesianRetrievalConfig:
    candidate_pool: int = 12
    likelihood_temperature: float = 2.0
    likelihood_floor: float = 0.02
    reliability_weight: float = 1.0
    authority_weight: float = 1.0
    freshness_weight: float = 0.8
    prior_floor: float = 1e-4
    min_relevance: float = 0.06
    min_top_posterior: float = 0.30
    min_query_coverage: float = 0.30
    min_dominance_ratio: float = 1.20


class BayesianIndex:
    """Generalized-Bayesian reranker over TF-IDF retrieval candidates.

    Candidate document hypotheses are scored by
        p(d|q) ∝ p(d) exp(tau * sim(q,d)).
    p(d) is induced by a Beta source-reliability mean plus authority/freshness
    metadata. The exponential similarity term is a Gibbs/generalized-Bayes
    likelihood, not claimed to be a literal language-model likelihood.
    """

    def __init__(self, documents: Sequence[DocumentRecord], config: BayesianRetrievalConfig | None = None):
        if not documents:
            raise ValueError("documents must be non-empty")
        self.documents = list(documents)
        self.config = config or BayesianRetrievalConfig()
        self.vectorizer = TfidfVectorizer(
            lowercase=True,
            stop_words="english",
            ngram_range=(1, 2),
            sublinear_tf=True,
            norm="l2",
        )
        self.matrix = self.vectorizer.fit_transform([d.retrieval_text for d in self.documents])

    def _log_prior(self, d: DocumentRecord) -> float:
        md = d.metadata
        a = max(float(md.get("reliability_alpha", 1.0)), 1e-6)
        b = max(float(md.get("reliability_beta", 1.0)), 1e-6)
        rel = _clip01(a / (a + b))
        auth = _clip01(md.get("authority", 0.5))
        fresh = _clip01(md.get("freshness", 0.5))
        cfg = self.config
        prior = max(
            cfg.prior_floor,
            (rel ** cfg.reliability_weight)
            * (auth ** cfg.authority_weight)
            * (fresh ** cfg.freshness_weight),
        )
        return math.log(prior)

    def lexical_scores(self, query: str) -> np.ndarray:
        q = self.vectorizer.transform([str(query)])
        # TF-IDF vectors are L2 normalized, so dot product is cosine similarity.
        return (self.matrix @ q.T).toarray().ravel()

    @staticmethod
    def _query_coverage(query: str, text: str) -> float:
        qterms = {t for t in re.findall(r"[A-Za-z0-9_-]+", query.lower()) if len(t) > 1 and t not in ENGLISH_STOP_WORDS}
        if not qterms:
            return 0.0
        dterms = set(re.findall(r"[A-Za-z0-9_-]+", text.lower()))
        return len(qterms & dterms) / len(qterms)

    def search(self, query: str, k: int = 4) -> list[RetrievalHit]:
        if k <= 0:
            return []
        scores = self.lexical_scores(query)
        pool_n = min(max(k, self.config.candidate_pool), len(self.documents))
        # Stable sort makes benchmark results reproducible under ties.
        cand = np.argsort(-scores, kind="stable")[:pool_n]
        logu: list[float] = []
        rows: list[tuple[int, float, float, float]] = []
        for i in cand:
            d = self.documents[int(i)]
            s = max(float(scores[int(i)]), 0.0)
            lp = self._log_prior(d)
            ll = self.config.likelihood_temperature * math.log(s + self.config.likelihood_floor)
            u = lp + ll
            rows.append((int(i), s, lp, ll))
            logu.append(u)
        lse = _logsumexp(logu)
        probs = [math.exp(u - lse) if math.isfinite(lse) else 1.0 / len(logu) for u in logu]
        ranked = sorted(
            zip(rows, probs),
            key=lambda rp: (rp[1], rp[0][1]),
            reverse=True,
        )
        hits: list[RetrievalHit] = []
        for rank, ((i, s, lp, ll), pr) in enumerate(ranked[:k], start=1):
            hits.append(RetrievalHit(
                document=self.documents[i],
                retrieval_score=s,
                log_prior=lp,
                log_likelihood=ll,
                log_posterior_unnormalized=lp + ll,
                posterior=float(pr),
                query_coverage=self._query_coverage(query, self.documents[i].retrieval_text),
                rank=rank,
            ))
        return hits

    def should_abstain(self, hits: Sequence[RetrievalHit]) -> bool:
        if not hits:
            return True
        top = hits[0]
        if (
            top.retrieval_score < self.config.min_relevance
            or top.posterior < self.config.min_top_posterior
            or top.query_coverage < self.config.min_query_coverage
        ):
            return True
        if len(hits) > 1:
            second = hits[1]
            dominance = top.posterior / max(second.posterior, 1e-12)
            if dominance < self.config.min_dominance_ratio:
                return True
        return False


class ExtractiveGroundedGenerator:
    """Dependency-free grounded answer synthesizer for tests and offline use.

    It selects the query-relevant sentence from the highest-posterior evidence
    and returns that evidence sentence verbatim with its document id. It never
    invents content not present in the retrieved context.
    """

    abstain_text = "I do not know based on the provided documents."

    @staticmethod
    def _sentences(text: str) -> list[str]:
        parts = re.split(r"(?<=[.!?])\s+|\n+", text.strip())
        return [p.strip() for p in parts if p.strip()]

    def generate(self, question: str, hits: Sequence[RetrievalHit]) -> tuple[str, list[str]]:
        if not hits:
            return self.abstain_text, []
        # The posterior ranking has already selected the evidence hypothesis.
        top = hits[0]
        sents = self._sentences(top.document.text)
        if not sents:
            return self.abstain_text, []
        if len(sents) == 1:
            sentence = sents[0]
        else:
            v = TfidfVectorizer(lowercase=True, stop_words="english", ngram_range=(1, 2))
            mat = v.fit_transform([question] + sents)
            q = mat[0]
            sims = (mat[1:] @ q.T).toarray().ravel()
            sentence = sents[int(np.argmax(sims))]
        return f"{sentence} [{top.document.id}]", [top.document.id]


class PlainTfidfIndex:
    """Conventional lexical baseline used by the end-to-end RAG benchmark."""
    def __init__(self, documents: Sequence[DocumentRecord]):
        self.documents = list(documents)
        self.vectorizer = TfidfVectorizer(lowercase=True, stop_words="english", ngram_range=(1, 2), sublinear_tf=True, norm="l2")
        self.matrix = self.vectorizer.fit_transform([d.retrieval_text for d in self.documents])

    @staticmethod
    def _query_coverage(query: str, text: str) -> float:
        qterms = {t for t in re.findall(r"[A-Za-z0-9_-]+", query.lower()) if len(t) > 1 and t not in ENGLISH_STOP_WORDS}
        if not qterms:
            return 0.0
        dterms = set(re.findall(r"[A-Za-z0-9_-]+", text.lower()))
        return len(qterms & dterms) / len(qterms)

    def search(self, query: str, k: int = 4) -> list[RetrievalHit]:
        q = self.vectorizer.transform([query])
        scores = (self.matrix @ q.T).toarray().ravel()
        order = np.argsort(-scores, kind="stable")[:k]
        # Baseline assigns normalized positive retrieval mass only for reporting.
        vals = [max(float(scores[int(i)]), 0.0) for i in order]
        z = sum(vals) or 1.0
        return [RetrievalHit(self.documents[int(i)], float(scores[int(i)]), 0.0, 0.0, 0.0, vals[j]/z, 0.0, j+1) for j, i in enumerate(order)]