from __future__ import annotations

from dataclasses import dataclass

from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
from .ports import ParentStore, QueryEmbeddingUnavailable, Retriever
from .sections import SectionResolver


@dataclass(frozen=True)
class EvidencePolicy:
    minimum_score: float = 0.12
    candidate_limit: int = 5
    evidence_limit: int = 3


class RetrievalService:
    """Section-filtered retrieval when the question names its attribute.

    Similarity is the fallback, not the default. Measured 2026-08-04, letting
    similarity choose the section answers "chống chỉ định" correctly 1 time in
    20, because the largest section (`duoc_ly_va_co_che_tac_dung`) sits close
    to any question about the drug. When the question says which section it
    wants, filtering answers it exactly.
    """

    def __init__(
        self,
        retriever: Retriever,
        parent_store: ParentStore,
        policy: EvidencePolicy | None = None,
        section_resolver: SectionResolver | None = None,
    ) -> None:
        self._retriever = retriever
        self._parent_store = parent_store
        self._policy = policy or EvidencePolicy()
        self._section_resolver = section_resolver

    def retrieve(self, query: str, drug_id: str) -> RetrievalResult:
        if not query.strip() or not drug_id.strip():
            return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_query_or_drug")

        section_hits = self._section_hits(query, drug_id)
        if section_hits is not None:
            # No `evidence_limit` here: the whole section is the answer, and a
            # truncated list of contraindications reads as a complete one.
            return self._decide(self._hydrate(section_hits, limit=None))

        # Drug resolved but no attribute named ("PARACETAMOL"): show the whole
        # monograph, in book order, rather than dead-ending on "specify an
        # attribute". A drug reference answers a drug name with the drug.
        overview_hits = self._drug_overview(drug_id)
        if overview_hits is not None:
            return self._decide(self._hydrate(overview_hits, limit=None))

        try:
            hits = self._retriever.search(
                query=query,
                drug_id=drug_id,
                limit=self._policy.candidate_limit,
            )
        except QueryEmbeddingUnavailable:
            # Fail closed. The section route needs no embedder, so this only
            # ever narrows the fallback: the caller is told nothing was found
            # rather than being shown an error page or, worse, an answer built
            # from a search that never ran.
            return RetrievalResult(
                EvidenceDecision.ABSTAIN, "query_embedding_unavailable"
            )
        if not hits or hits[0].score < self._policy.minimum_score:
            return RetrievalResult(EvidenceDecision.ABSTAIN, "insufficient_retrieval_score")
        return self._decide(self._hydrate(hits))

    def _drug_overview(self, drug_id: str) -> list[SearchHit] | None:
        """Every prose section of the drug, or None if the store cannot scroll."""
        find_by_drug = getattr(self._retriever, "find_by_drug", None)
        if find_by_drug is None:
            return None
        hits = find_by_drug(drug_id)
        return hits or None

    def _section_hits(self, query: str, drug_id: str) -> list[SearchHit] | None:
        """Hits for an explicitly named section, or None to fall back.

        Returns None — not an empty list — when this route does not apply, so
        "no section named" stays distinguishable from "section named but empty".
        """
        if self._section_resolver is None:
            return None
        find_by_section = getattr(self._retriever, "find_by_section", None)
        if find_by_section is None:
            return None
        match = self._section_resolver.resolve(query)
        if match is None:
            return None
        hits = find_by_section(drug_id, match.section_key)
        return hits or None

    def _decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult:
        if not evidence:
            return RetrievalResult(EvidenceDecision.ABSTAIN, "parent_hydration_failed")
        if any(not item.source_refs for item in evidence):
            return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_provenance")
        if any(item.requires_visual_check for item in evidence):
            return RetrievalResult(EvidenceDecision.VERIFY_PDF, "visual_verification_required", evidence)
        return RetrievalResult(EvidenceDecision.ANSWERABLE, "grounded_evidence_available", evidence)

    def _hydrate(
        self, hits: list[SearchHit], limit: int | None = -1
    ) -> tuple[Evidence, ...]:
        output: list[Evidence] = []
        seen: set[str] = set()
        for hit in hits:
            document = hit.document
            evidence_id = document.parent_id or document.doc_id
            if evidence_id in seen:
                continue
            seen.add(evidence_id)
            if document.parent_id:
                parent = self._parent_store.get(document.parent_id)
                if parent is None:
                    continue
                output.append(Evidence(
                    evidence_id=parent.parent_id,
                    matched_doc_id=document.doc_id,
                    kind=parent.kind,
                    text=parent.text,
                    score=hit.score,
                    source_refs=parent.source_refs,
                    hydrated_from_parent=True,
                    requires_visual_check=(
                        document.requires_visual_check or parent.requires_visual_check
                    ),
                ))
            else:
                output.append(Evidence(
                    evidence_id=document.doc_id,
                    matched_doc_id=document.doc_id,
                    kind=document.kind,
                    text=document.text,
                    score=hit.score,
                    source_refs=document.source_refs,
                    hydrated_from_parent=False,
                    requires_visual_check=document.requires_visual_check,
                ))
            cap = self._policy.evidence_limit if limit == -1 else limit
            if cap is not None and len(output) >= cap:
                break
        return tuple(output)