"""Section routing: the fix for hit@1 0.05 on `chong_chi_dinh`.

Measured 2026-08-04 on the real Cohere collection, letting vector similarity
choose the section answered contraindication questions correctly 1 time in 20.
These tests pin the two properties that make filtering safe: the longer phrase
always wins, and an unrecognised question routes nowhere rather than guessing.
"""
from __future__ import annotations

from rag.models import (
    EvidenceDecision,
    RetrievalDocument,
    SearchHit,
    SourceRef,
)
from rag.in_memory import InMemoryParentStore
from rag.sections import SectionResolver
from rag.service import EvidencePolicy, RetrievalService

SOURCE = SourceRef(physical_page=200, precision="page", printed_page=142)


def _doc(doc_id: str, section_key: str, text: str) -> RetrievalDocument:
    return RetrievalDocument(
        doc_id=doc_id,
        parent_id=None,
        drug_id="aspirin",
        kind="prose",
        section_key=section_key,
        text=text,
        source_refs=(SOURCE,),
        requires_visual_check=False,
    )


class SectionAwareRetriever:
    """Fake that records which route the service actually took."""

    def __init__(self, docs: list[RetrievalDocument]) -> None:
        self._docs = docs
        self.search_calls: list[str] = []
        self.section_calls: list[tuple[str, str]] = []

    def search(
        self, query: str, drug_id: str, limit: int  # noqa: ARG002 — Retriever protocol
    ) -> list[SearchHit]:
        self.search_calls.append(query)
        # Deliberately wrong on purpose: the whole point is that the section
        # route must not consult similarity at all.
        return [SearchHit(self._docs[-1], 0.99)]

    def find_by_section(self, drug_id: str, section_key: str) -> list[SearchHit]:
        self.section_calls.append((drug_id, section_key))
        return [
            SearchHit(doc, 1.0) for doc in self._docs if doc.section_key == section_key
        ]


class LexicalAwareRetriever(SectionAwareRetriever):
    """Adds `search_lexical`, scripted per test rather than doing real
    text matching — this suite is about `_pooled_neighbour_hits`' bounding
    and threshold logic, not the lexical scorer itself (see
    `adapters/qdrant.py`'s own coverage for that)."""

    def __init__(self, docs: list[RetrievalDocument], lexical_hits: list[SearchHit]) -> None:
        super().__init__(docs)
        self._lexical_hits = lexical_hits
        self.lexical_calls: list[tuple[str, str]] = []

    def search_lexical(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:  # noqa: ARG002
        self.lexical_calls.append((query, drug_id))
        return self._lexical_hits


class SimilarityOnlyRetriever:
    def __init__(self, docs: list[RetrievalDocument]) -> None:
        self._docs = docs
        self.search_calls: list[str] = []

    def search(
        self, query: str, drug_id: str, limit: int  # noqa: ARG002 — Retriever protocol
    ) -> list[SearchHit]:
        self.search_calls.append(query)
        return [SearchHit(self._docs[0], 0.99)]


CONTRA = [
    _doc("c1", "chong_chi_dinh", "Mẫn cảm với aspirin."),
    _doc("c2", "chong_chi_dinh", "Loét dạ dày tá tràng đang tiến triển."),
    _doc("c3", "chong_chi_dinh", "Hen do aspirin."),
    _doc("c4", "chong_chi_dinh", "Suy gan nặng."),
    _doc("c5", "chong_chi_dinh", "Trẻ em dưới 16 tuổi có sốt virus."),
]
INDICATION = [_doc("i1", "chi_dinh", "Giảm đau, hạ sốt, chống viêm.")]
PHARMACOLOGY = [_doc("p1", "duoc_ly_va_co_che_tac_dung", "Ức chế cyclooxygenase.")]
PRECAUTION = [_doc("t1", "than_trong", "Thận trọng với người suy thận.")]
INTERACTION = [_doc("x1", "tuong_tac_thuoc", "Tương tác với thuốc chống đông máu.")]
ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY + PRECAUTION + INTERACTION


def _service(retriever, resolver: SectionResolver | None) -> RetrievalService:
    return RetrievalService(
        retriever,
        InMemoryParentStore([]),
        EvidencePolicy(minimum_score=0.01),
        section_resolver=resolver,
    )


class TestSectionResolver:
    def test_contraindication_is_never_read_as_indication(self) -> None:
        """The one that measured 0.05. "chống chỉ định" contains "chỉ định"."""
        resolver = SectionResolver()
        assert resolver.resolve("Chống chỉ định của aspirin là gì?").section_key == (
            "chong_chi_dinh"
        )
        assert resolver.resolve("Chỉ định của aspirin?").section_key == "chi_dinh"

    def test_works_without_diacritics(self) -> None:
        assert SectionResolver().resolve("aspirin chong chi dinh").section_key == (
            "chong_chi_dinh"
        )

    def test_overdose_is_not_read_as_dose(self) -> None:
        resolver = SectionResolver()
        assert resolver.resolve("xử trí quá liều metformin").section_key == (
            "qua_lieu_va_xu_tri"
        )
        assert resolver.resolve("liều dùng metformin").section_key == (
            "lieu_luong_va_cach_dung"
        )

    def test_adr_management_is_not_read_as_adr_itself(self) -> None:
        resolver = SectionResolver()
        assert resolver.resolve("xử trí tác dụng phụ của prednisolon").section_key == (
            "huong_dan_xu_tri_adr"
        )
        assert resolver.resolve("tác dụng phụ của prednisolon").section_key == (
            "tac_dung_khong_mong_muon"
        )

    def test_incompatibility_is_not_read_as_interaction(self) -> None:
        resolver = SectionResolver()
        assert resolver.resolve("tương kỵ của ceftriaxon").section_key == "tuong_ky"
        assert resolver.resolve("tương tác của ceftriaxon").section_key == (
            "tuong_tac_thuoc"
        )

    def test_bare_lieu_resolves_without_capturing_overdose(self) -> None:
        """Found by testing on human-written golden questions, not templates.

        4 of 16 said just "Liều Metformin cho người lớn?". Adding bare "liều"
        is only safe because "quá liều" is longer and is tested first.
        """
        resolver = SectionResolver()
        assert resolver.resolve("Liều Metformin cho người lớn?").section_key == (
            "lieu_luong_va_cach_dung"
        )
        assert resolver.resolve("quá liều paracetamol").section_key == (
            "qua_lieu_va_xu_tri"
        )

    def test_colloquial_pregnancy_phrasing(self) -> None:
        assert SectionResolver().resolve(
            "Bà bầu dùng Ibuprofen được không?"
        ).section_key == "thoi_ky_mang_thai"

    def test_unrecognised_question_routes_nowhere(self) -> None:
        """No match must not become a guess — the caller falls back."""
        assert SectionResolver().resolve("thuốc này giá bao nhiêu") is None
        assert SectionResolver().resolve("") is None

    def test_new_section_needs_no_code_change(self) -> None:
        resolver = SectionResolver({"invented_section": ("một mục hoàn toàn mới",)})
        assert resolver.resolve("hỏi về một mục hoàn toàn mới").section_key == (
            "invented_section"
        )


class TestSectionRouting:
    def test_named_section_bypasses_similarity_entirely(self) -> None:
        retriever = SectionAwareRetriever(ALL_DOCS)
        result = _service(retriever, SectionResolver()).retrieve(
            "Chống chỉ định của aspirin là gì?", "aspirin"
        )
        assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
        assert retriever.search_calls == []
        assert result.decision == EvidenceDecision.ANSWERABLE
        assert {item.evidence_id for item in result.evidence} == {
            "c1", "c2", "c3", "c4", "c5",
        }

    def test_whole_section_is_returned_past_the_evidence_limit(self) -> None:
        """Five contraindications must not arrive as three."""
        retriever = SectionAwareRetriever(ALL_DOCS)
        service = RetrievalService(
            retriever,
            InMemoryParentStore([]),
            EvidencePolicy(minimum_score=0.01, evidence_limit=3),
            section_resolver=SectionResolver(),
        )
        result = service.retrieve("chống chỉ định aspirin", "aspirin")
        assert len(result.evidence) == 5

    def test_unnamed_section_falls_back_to_similarity(self) -> None:
        retriever = SectionAwareRetriever(ALL_DOCS)
        result = _service(retriever, SectionResolver()).retrieve(
            "aspirin dùng cho bệnh nhân này thế nào", "aspirin"
        )
        assert retriever.section_calls == []
        assert retriever.search_calls
        assert result.decision == EvidenceDecision.ANSWERABLE

    def test_retriever_without_the_capability_still_works(self) -> None:
        retriever = SimilarityOnlyRetriever(ALL_DOCS)
        result = _service(retriever, SectionResolver()).retrieve(
            "Chống chỉ định của aspirin là gì?", "aspirin"
        )
        assert retriever.search_calls
        assert result.decision == EvidenceDecision.ANSWERABLE

    def test_no_resolver_keeps_the_old_behaviour(self) -> None:
        retriever = SectionAwareRetriever(ALL_DOCS)
        _service(retriever, None).retrieve("chống chỉ định aspirin", "aspirin")
        assert retriever.section_calls == []
        assert retriever.search_calls

    def test_lexically_strong_neighbour_section_is_pooled_in(self) -> None:
        """A precaution some drug's own "thận trọng" text never mentions can
        still be filed under "chống chỉ định" (found live: Aspirin + loét dạ
        dày). A neighbour section with real term overlap (score above
        threshold) gets pooled in whole, not just the matching fragment."""
        retriever = LexicalAwareRetriever(
            ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 6.0)],  # c2: "chong_chi_dinh"
        )
        result = _service(retriever, SectionResolver()).retrieve(
            "Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"
        )
        assert retriever.section_calls == [
            ("aspirin", "than_trong"), ("aspirin", "chong_chi_dinh"),
        ]
        assert retriever.lexical_calls == [
            ("Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"),
        ]
        # Whole section pooled (c1-c5), not just the one matching chunk (c2).
        assert {item.evidence_id for item in result.evidence} == {
            "t1", "c1", "c2", "c3", "c4", "c5",
        }

    def test_weak_lexical_match_is_not_pooled(self) -> None:
        """One incidental token overlap must not drag in a whole tangential
        section — only real term overlap (>= threshold) counts."""
        retriever = LexicalAwareRetriever(
            ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 4.0)],
        )
        result = _service(retriever, SectionResolver()).retrieve(
            "Thận trọng khi dùng aspirin là gì?", "aspirin"
        )
        assert retriever.section_calls == [("aspirin", "than_trong")]
        assert {item.evidence_id for item in result.evidence} == {"t1"}

    def test_lexical_pooling_is_bounded(self) -> None:
        """At most 2 neighbour sections pool in, even if more score above
        threshold — unbounded pooling is precision loss dressed as recall."""
        retriever = LexicalAwareRetriever(
            ALL_DOCS,
            lexical_hits=[
                SearchHit(CONTRA[0], 7.0),      # chong_chi_dinh
                SearchHit(INDICATION[0], 6.0),  # chi_dinh
                SearchHit(INTERACTION[0], 5.0), # tuong_tac_thuoc — 3rd, over bound
            ],
        )
        result = _service(retriever, SectionResolver()).retrieve(
            "Thận trọng khi dùng aspirin là gì?", "aspirin"
        )
        assert retriever.section_calls == [
            ("aspirin", "than_trong"),
            ("aspirin", "chong_chi_dinh"),
            ("aspirin", "chi_dinh"),
        ]
        # tuong_tac_thuoc (3rd-ranked lexical hit) never pooled — the bound
        # (2 neighbours) already used by chong_chi_dinh + chi_dinh.
        assert "x1" not in {item.evidence_id for item in result.evidence}

    def test_generic_pharmacology_section_is_never_pooled(self) -> None:
        """`duoc_ly_va_co_che_tac_dung` is the corpus's documented
        false-positive attractor (see `sections.py`) — excluded outright,
        even with a lexical score that would otherwise clear the threshold
        (measured live: it scored a close second right behind the real
        answer on the exact query that motivated this whole mechanism)."""
        retriever = LexicalAwareRetriever(
            ALL_DOCS, lexical_hits=[SearchHit(PHARMACOLOGY[0], 9.0)],
        )
        result = _service(retriever, SectionResolver()).retrieve(
            "Thận trọng khi dùng aspirin là gì?", "aspirin"
        )
        assert retriever.section_calls == [("aspirin", "than_trong")]
        assert "p1" not in {item.evidence_id for item in result.evidence}

    def test_retriever_without_lexical_search_still_works(self) -> None:
        """No `search_lexical` on the retriever (interface segregation, same
        pattern as `find_by_section`/`find_by_drug`) — pooling is skipped,
        not an error."""
        retriever = SectionAwareRetriever(ALL_DOCS)
        result = _service(retriever, SectionResolver()).retrieve(
            "Thận trọng khi dùng aspirin là gì?", "aspirin"
        )
        assert retriever.section_calls == [("aspirin", "than_trong")]
        assert {item.evidence_id for item in result.evidence} == {"t1"}

    def test_named_but_empty_section_falls_back(self) -> None:
        """A drug with no such section must not abstain — similarity still tries."""
        retriever = SectionAwareRetriever(INDICATION + PHARMACOLOGY)
        result = _service(retriever, SectionResolver()).retrieve(
            "chống chỉ định aspirin", "aspirin"
        )
        assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
        assert retriever.search_calls
        assert result.decision == EvidenceDecision.ANSWERABLE