stack/tests/llm/test_evidence.py

"""llm.evidence — valuation prompt block, SSE payload, code-cited rule sources."""

from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import duckdb
import pytest

from llm import evidence
from llm.config import LlmConfig
from llm.evidence import (
    ValuationEvidence,
    code_cited_sources,
    manual_sources,
    merge_sources,
    valuation_evidence,
    warm,
)
from pfs.codetables import ensure_tables
from pfs.valuation import ValuationRow

CFG = LlmConfig(
    ollama_hosts=("http://h1:11434",),
    embed_model="e",
    instruct_model="c",
    embed_dim=768,
    build_ann_index=False,
    pg_host="x",
    pg_port=5432,
    pg_db="llm",
    pg_user="llm",
    duckdb_replica="/nonexistent/aco.ro.duckdb",
    valuation_years=4,
    code_cited_per_code=2,
)


@pytest.fixture(autouse=True)
def _clear_replica_cache():
    """The replica handle is process-global; no test may inherit another's."""
    evidence._REPLICA = None
    yield
    evidence._REPLICA = None


def _row(**kw) -> ValuationRow:
    base = dict(
        code="G0556",
        description="Adv prim care mgmt lvl 1",
        vintage="CY2026 final",
        year=2026,
        proposed=False,
        status="A",
        work=0.25,
        pe_nf=0.22,
        pe_f=0.06,
        mp=0.02,
        total_nf=0.49,
        total_f=0.33,
        cf=33.4009,
        pay_nf=16.37,
        pay_f=11.02,
        label="[PFS CY2026 Addendum B]",
        citation="90 FR 49266",
        url="https://www.federalregister.gov/citation/90-FR-49266",
    )
    base.update(kw)
    return ValuationRow(**base)


class TestEvidenceText:
    def test_prompt_block_exact(self):
        ev = ValuationEvidence(
            codes=("G0556",), families=("APCM",), rows=(_row(),), unpriced=("Z9999",)
        )
        assert ev.prompt_block() == (
            "Valuation (authoritative for RVUs, conversion factors and payments; "
            "national unadjusted, no GPCI):\n"
            '[PFS CY2026 Addendum B] G0556 "Adv prim care mgmt lvl 1" — status A; work 0.25; '
            "PE 0.22 non-fac / 0.06 fac; MP 0.02; total 0.49 non-fac / 0.33 fac; "
            "CF $33.4009 → payment $16.37 non-fac / $11.02 fac\n"
            "Not priced in the indexed fee schedule: Z9999"
        )

    def test_prompt_block_labels_the_cf_and_an_unpaid_status(self):
        ev = ValuationEvidence(
            ("G0559",),
            (),
            (
                _row(
                    code="G0559",
                    status="B",
                    status_note="bundled — no separate PFS payment",
                    cf_note="non-APM standard CF; QP CF $33.5675",
                    pay_nf=None,
                    pay_f=None,
                ),
            ),
            (),
        )
        line = ev.prompt_block().splitlines()[1]
        assert "— status B; bundled — no separate PFS payment; work" in line
        assert "CF $33.4009 [non-APM standard CF; QP CF $33.5675] → payment" in line
        assert "payment n/a non-fac / n/a fac" in line

    def test_prompt_block_null_components(self):
        ev = ValuationEvidence(
            ("G9999",),
            (),
            (
                _row(
                    code="G9999",
                    status="I",
                    work=None,
                    pe_nf=None,
                    pe_f=None,
                    mp=None,
                    total_nf=None,
                    total_f=None,
                    pay_nf=None,
                    pay_f=None,
                ),
            ),
            (),
        )
        line = ev.prompt_block().splitlines()[1]
        assert (
            "status I; work n/a; PE n/a non-fac / n/a fac" in line
            and "payment n/a" in line
        )

    def test_payload_shape(self):
        ev = ValuationEvidence(
            ("G0556",),
            ("APCM",),
            (
                _row(),
                _row(
                    year=2027,
                    proposed=True,
                    vintage="CY2027 proposed",
                    label="[CY2027 NPRM Addendum B]",
                    citation="91 FR 43842",
                    url="https://www.federalregister.gov/citation/91-FR-43842",
                ),
            ),
            (),
        )
        p = ev.payload()
        assert (
            p["type"] == "valuation"
            and p["codes"] == ["G0556"]
            and p["families"] == ["APCM"]
        )
        assert p["rows"][0]["pay_nf"] == 16.37 and p["rows"][1]["proposed"] is True
        assert p["rows"][0]["status_note"] == "" and p["rows"][0]["cf_note"] == ""
        assert p["provenance"] == [
            {
                "label": "[PFS CY2026 Addendum B]",
                "vintage": "CY2026 final",
                "citation": "90 FR 49266",
                "url": "https://www.federalregister.gov/citation/90-FR-49266",
            },
            {
                "label": "[CY2027 NPRM Addendum B]",
                "vintage": "CY2027 proposed",
                "citation": "91 FR 43842",
                "url": "https://www.federalregister.gov/citation/91-FR-43842",
            },
        ]
        assert p["unpriced"] == []


class TestReplicaPath:
    def test_absolute_path_is_returned_unchanged(self):
        cfg = replace(CFG, duckdb_replica="/nonexistent/aco.ro.duckdb")
        assert evidence._replica_path(cfg) == "/nonexistent/aco.ro.duckdb"

    def test_relative_path_is_resolved_against_conf_root(self):
        from conf import ROOT

        cfg = replace(CFG, duckdb_replica="data/replica/aco.ro.duckdb")
        assert evidence._replica_path(cfg) == str(ROOT / "data/replica/aco.ro.duckdb")


class TestValuationEvidence:
    def test_none_when_no_codes(self):
        assert (
            valuation_evidence("what did commenters say about telehealth?", CFG) is None
        )

    @patch("llm.evidence.valuation", return_value=([_row()], []))
    @patch("llm.evidence.duckdb.connect")
    def test_detects_and_queries_replica_read_only(self, mock_connect, mock_val):
        ev = valuation_evidence("How is APCM valued?", CFG)
        mock_connect.assert_called_once_with(
            "/nonexistent/aco.ro.duckdb", read_only=True
        )
        assert mock_val.call_args.args[1] == ["G0556", "G0557", "G0558"]
        assert mock_val.call_args.kwargs == {"years": 4}
        assert (
            ev is not None
            and ev.codes == ("G0556", "G0557", "G0558")
            and ev.families == ("APCM",)
        )
        mock_connect.return_value.close.assert_not_called()  # cached, not closed

    @patch("llm.evidence.duckdb.connect", side_effect=OSError("no replica"))
    def test_unopenable_replica_yields_none(self, _c, caplog):
        assert valuation_evidence("G0556?", CFG) is None
        assert "valuation evidence skipped" in caplog.text

    @patch("llm.evidence.valuation", return_value=([_row()], []))
    @patch("llm.evidence.duckdb.connect")
    def test_replica_connection_is_reused(self, mock_connect, _val):
        valuation_evidence("G0556?", CFG)
        valuation_evidence("G0557?", CFG)
        mock_connect.assert_called_once()
        mock_connect.return_value.close.assert_not_called()

    @patch("llm.evidence.valuation", return_value=([_row()], []))
    @patch("llm.evidence.duckdb.connect")
    def test_republished_replica_reopens_and_closes_the_old_handle(
        self, mock_connect, _val, monkeypatch
    ):
        first, second = MagicMock(), MagicMock()
        mock_connect.side_effect = [first, second]
        mtimes = iter([1, 2])
        monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
        valuation_evidence("G0556?", CFG)
        valuation_evidence("G0557?", CFG)
        assert mock_connect.call_count == 2
        first.close.assert_called_once()
        second.close.assert_not_called()

    @patch("llm.evidence.valuation", return_value=([], []))
    @patch("llm.evidence.duckdb.connect")
    def test_many_codes_shorten_the_window(self, _c, mock_val):
        valuation_evidence("How is CCM valued?", CFG)  # six-code family
        assert mock_val.call_args.kwargs == {"years": 2}

    @patch("llm.evidence.valuation", side_effect=RuntimeError("query boom"))
    @patch("llm.evidence.duckdb.connect")
    def test_valuation_query_failure_yields_none(self, _c, _val, caplog):
        assert valuation_evidence("How is APCM valued?", CFG) is None
        assert "valuation evidence skipped (query)" in caplog.text


class TestCapCodes:
    """``cap_codes`` — Ruling B11 per-turn code cap."""

    def test_noop_at_or_under_the_cap(self):
        from pfs.families import Detection

        det = Detection(codes=("A", "B"), families=(), explicit=("A",), wide=())
        assert evidence.cap_codes(det, 2) is det.codes

    def test_explicit_codes_survive_first(self):
        from pfs.families import Detection

        det = Detection(
            codes=tuple(f"{i:05d}" for i in range(10)),
            families=(),
            explicit=("00009", "00005"),
            wide=(),
        )
        capped = evidence.cap_codes(det, 3)
        assert len(capped) == 3
        assert {"00009", "00005"} <= set(capped)

    def test_remaining_slots_fill_in_family_order(self, monkeypatch):
        from pfs.families import Detection, Family

        fam_a = Family("FAMA", "Family A", ("A1", "A2"), ())
        fam_b = Family("FAMB", "Family B", ("B1", "B2"), ())
        monkeypatch.setattr(evidence, "FAMILIES", {"FAMA": fam_a, "FAMB": fam_b})
        det = Detection(
            codes=("A1", "A2", "B1", "B2"),
            families=("FAMB", "FAMA"),  # detected order — B before A
            explicit=(),
            wide=(),
        )
        capped = evidence.cap_codes(det, 3)
        assert capped == ("B1", "B2", "A1")

    def test_logs_the_drop_count(self, caplog):
        from pfs.families import Detection

        det = Detection(
            codes=tuple(f"{i:05d}" for i in range(5)),
            families=(),
            explicit=(),
            wide=(),
        )
        evidence.cap_codes(det, 2)
        assert "dropped 3 of 5" in caplog.text

    def test_unregistered_family_key_is_skipped(self, monkeypatch):
        from pfs.families import Detection, Family

        fam_a = Family("FAMA", "Family A", ("A1", "A2"), ())
        # FAMB is in det.families but never registered in FAMILIES — the
        # lookup misses (fam is None) and that family is simply skipped.
        monkeypatch.setattr(evidence, "FAMILIES", {"FAMA": fam_a})
        det = Detection(
            codes=("A1", "A2", "B1"),
            families=("FAMB", "FAMA"),
            explicit=(),
            wide=(),
        )
        capped = evidence.cap_codes(det, 2)
        assert capped == ("A1", "A2")


class TestValuationRowCap:
    """Ruling B11: ``ValuationEvidence.prompt_block`` renders only the
    capped rows (explicit codes first, then newest vintage); ``payload``
    always carries every row."""

    def _rows(self):
        return [
            _row(code="G0556", year=2023, label="[Y2023]"),
            _row(code="G0556", year=2024, label="[Y2024]"),
            _row(code="G0557", year=2025, label="[Y2025]"),
            _row(code="G0558", year=2026, label="[Y2026]"),
        ]

    def test_payload_keeps_every_row_regardless_of_cap(self):
        ev = ValuationEvidence(
            ("G0556", "G0557", "G0558"), (), tuple(self._rows()), (), max_prompt_rows=2
        )
        assert len(ev.payload()["rows"]) == 4

    def test_prompt_block_caps_and_prefers_explicit_then_newest_vintage(self):
        ev = ValuationEvidence(
            ("G0556", "G0557", "G0558"),
            (),
            tuple(self._rows()),
            (),
            explicit=("G0557",),
            max_prompt_rows=2,
        )
        block = ev.prompt_block()
        assert "[Y2025]" in block  # explicit code's row survives
        assert "[Y2026]" in block  # newest vintage among the rest
        assert "[Y2023]" not in block and "[Y2024]" not in block

    def test_max_rows_override_for_one_call(self):
        ev = ValuationEvidence(
            ("G0556", "G0557", "G0558"), (), tuple(self._rows()), (), max_prompt_rows=24
        )
        assert ev.prompt_block().count("\n") + 1 == 5  # header + 4 rows, no cap
        assert ev.prompt_block(max_rows=1).count("[Y") == 1


class TestWarm:
    """#699 ruling B3: warm the replica connection ahead of the first
    real request — at API startup and again (a no-op once cached) at the
    top of every ``valuation_evidence`` call."""

    def test_warm_opens_once_and_is_a_noop_on_replay(self, tmp_path):
        path = tmp_path / "aco.ro.duckdb"
        path.touch()  # warm() only needs the file to exist, not be a real db
        cfg = replace(CFG, duckdb_replica=str(path))
        with patch("llm.evidence.duckdb.connect") as mock_connect:
            mock_connect.return_value = MagicMock()
            warm(cfg)
            warm(cfg)
        mock_connect.assert_called_once_with(str(path), read_only=True)

    def test_warm_is_a_noop_when_the_replica_file_is_missing(self):
        with patch("llm.evidence.duckdb.connect") as mock_connect:
            warm(CFG)  # CFG.duckdb_replica ("/nonexistent/...") doesn't exist
        mock_connect.assert_not_called()

    def test_warm_swallows_connect_failures(self, tmp_path, caplog):
        path = tmp_path / "aco.ro.duckdb"
        path.touch()
        cfg = replace(CFG, duckdb_replica=str(path))
        with patch("llm.evidence.duckdb.connect", side_effect=OSError("boom")):
            warm(cfg)  # must not raise
        assert "replica warm-up skipped" in caplog.text

    @patch("llm.evidence.detect_codes")
    @patch("llm.evidence.warm")
    def test_valuation_evidence_warms_before_detecting_codes(
        self, mock_warm, mock_detect
    ):
        order: list[str] = []
        mock_warm.side_effect = lambda cfg: order.append("warm")

        def _detect(question):
            order.append("detect")
            return SimpleNamespace(codes=())

        mock_detect.side_effect = _detect
        assert valuation_evidence("anything", CFG) is None
        assert order == ["warm", "detect"]
        mock_warm.assert_called_once_with(CFG)


class TestConnectRefreshesFamilies:
    """#699: derived families (``pfs.code_family``) live only on the
    replica — ``_connect`` is the chat's one hook to pick them up, on
    every open and every reopen (a republished replica)."""

    @patch("llm.evidence.refresh_from")
    @patch("llm.evidence.duckdb.connect")
    def test_connect_refreshes_families_on_open_and_reopen(
        self, mock_connect, mock_refresh, monkeypatch
    ):
        mtimes = iter([1, 1, 2])
        monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
        mock_connect.side_effect = [MagicMock(), MagicMock()]

        con1 = evidence._connect("/x")
        assert mock_refresh.call_count == 1
        mock_refresh.assert_called_with(con1)

        con2 = evidence._connect("/x")  # same (path, mtime) — cached, no refresh
        assert con2 is con1
        assert mock_refresh.call_count == 1

        con3 = evidence._connect("/x")  # mtime changed — reopen, refresh again
        assert mock_refresh.call_count == 2
        mock_refresh.assert_called_with(con3)

    @patch("llm.evidence.refresh_from", side_effect=RuntimeError("boom"))
    @patch("llm.evidence.duckdb.connect")
    def test_refresh_failure_does_not_break_the_handle(
        self, _mock_connect, _mock_refresh, caplog
    ):
        con = evidence._connect("/x")
        assert con is not None
        assert "family refresh skipped" in caplog.text

    @patch("llm.evidence.duckdb.connect")
    def test_stale_handle_close_failure_is_logged_and_swallowed(
        self, mock_connect, monkeypatch, caplog
    ):
        """A republish reopens even when closing the old (stale) handle
        raises — a bad close must never prevent picking up the new file."""
        first, second = MagicMock(), MagicMock()
        first.close.side_effect = RuntimeError("close boom")
        mock_connect.side_effect = [first, second]
        mtimes = iter([1, 2])
        monkeypatch.setattr(evidence, "_mtime", lambda _p: next(mtimes))
        con1 = evidence._connect("/x")
        con2 = evidence._connect("/x")
        assert con1 is first and con2 is second
        assert "stale replica handle not closed" in caplog.text


RVU_COLS = (
    "hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
    "work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
    "non_fac_total DOUBLE, fac_total DOUBLE, conv_factor DOUBLE, year INTEGER"
)
PROPOSED_COLS = (
    "hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
    "work_rvu DOUBLE, non_fac_pe_rvu DOUBLE, fac_pe_rvu DOUBLE, mp_rvu DOUBLE, "
    "cms_rule_id VARCHAR"
)


class TestConcurrentChats:
    """/chat streams run in Starlette's threadpool, so several turns share
    the cached replica handle. Without a per-call cursor their result sets
    cross-talk and the broad except swallows it as "no valuation"."""

    @pytest.fixture
    def replica(self, tmp_path):
        db = tmp_path / "aco.ro.duckdb"
        con = duckdb.connect(str(db))
        con.execute("CREATE SCHEMA pfs")
        con.execute(f"CREATE TABLE pfs.rvu ({RVU_COLS})")
        con.execute(f"CREATE TABLE pfs.rvu_proposed ({PROPOSED_COLS})")
        # two unfamilied codes with different row counts, so a swapped
        # result set is visible as a wrong count
        con.executemany(
            "INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
            [
                (
                    "99213",
                    None,
                    "Office visit",
                    "A",
                    1.0,
                    1.0,
                    0.5,
                    0.05,
                    2.05,
                    1.55,
                    30.0,
                    2026,
                ),
                *[
                    (
                        "99214",
                        None,
                        "Office visit",
                        "A",
                        1.5,
                        1.0,
                        0.5,
                        0.05,
                        2.55,
                        2.05,
                        30.0,
                        y,
                    )
                    for y in (2024, 2025, 2026)
                ],
            ],
        )
        con.close()
        return db

    def test_parallel_turns_share_one_handle_without_cross_talk(self, replica):
        cfg = replace(CFG, duckdb_replica=str(replica))
        real_connect, opens = duckdb.connect, []

        def counting_connect(*a, **kw):
            opens.append(a)
            return real_connect(*a, **kw)

        expected = {"99213": 1, "99214": 3}
        questions = ["What does 99213 pay?", "What does 99214 pay?"]

        def one(i):
            q = questions[i % 2]
            ev = valuation_evidence(q, cfg)
            assert ev is not None, "valuation dropped under concurrency"
            (code,) = ev.codes
            return code, len(ev.rows)

        with patch("llm.evidence.duckdb.connect", counting_connect):
            with ThreadPoolExecutor(max_workers=8) as pool:
                results = list(pool.map(one, range(160)))

        assert all(n == expected[code] for code, n in results), sorted(set(results))
        assert len(opens) == 1  # one handle, shared

    def test_cursor_is_closed_but_the_handle_is_kept(self, replica):
        cfg = replace(CFG, duckdb_replica=str(replica))
        assert valuation_evidence("99213?", cfg) is not None
        assert evidence._REPLICA is not None
        # the cached parent still answers — it was never closed
        assert valuation_evidence("99214?", cfg) is not None


def _engine(rows):
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    conn.execute.return_value.fetchall.return_value = rows
    return engine, conn


def _rule_md(p_id, codes, date="2024-12-09"):
    return {
        "kind": "rule",
        "item_key": "R1",
        "p_id": str(p_id),
        "page": "97710",
        "ordinal": "1",
        "html_url": "https://www.federalregister.gov/d/2024-25382",
        "fr_volume": "89",
        "date": date,
        "title": "CY2025 PFS final rule",
        "codes": codes,
    }


def _comment_md(seq, codes, item_key="C1", date="2024-12-09"):
    return {
        "kind": "comment",
        "item_key": item_key,
        "seq": str(seq),
        "comment_id": "CMS-2026-2377-3438",
        "date": date,
        "codes": codes,
        "families": "",
    }


def _corpus_md(seq, codes="", families="", item_key="CPT-ED", date="2024-12-09"):
    return {
        "kind": "corpus",
        "item_key": item_key,
        "seq": str(seq),
        "url": "https://example.org/cpt.pdf",
        "title": "CPT Manual",
        "date": date,
        "codes": codes,
        "families": families,
    }


def _multi_engine(by_code=None, by_family=None):
    """A fake engine whose rows depend on which collection *and* which of
    the two SQL statements (codes vs. families) is being executed —
    distinguished by which bind parameter the call carries."""
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    by_code = by_code or {}
    by_family = by_family or {}

    def _execute(_sql, params):
        collection = params["collection"]
        table = by_code if "codes" in params else by_family
        result = MagicMock()
        result.fetchall.return_value = table.get(collection, [])
        return result

    conn.execute.side_effect = _execute
    return engine, conn


class TestCodeCitedSources:
    def test_sql_and_per_code_cap_and_dedupe(self):
        engine, conn = _engine(
            [
                ("para one G0556", _rule_md(1, "G0556")),
                ("para two G0556", _rule_md(2, "G0556")),
                ("para three G0556", _rule_md(3, "G0556")),  # over the cap for G0556
                (
                    "para four G0557",
                    _rule_md(4, "G0557 G0556"),
                ),  # G0557 still under cap
                ("para one again", _rule_md(1, "G0556")),  # duplicate (item_key, p_id)
            ]
        )
        out = code_cited_sources(engine, ["G0556", "G0557"], per_code=2)
        sql = str(conn.execute.call_args.args[0])
        assert "string_to_array" in sql and "langchain_pg_collection" in sql
        assert "PARTITION BY" in sql
        # the collection is a scalar subquery, not a join to filter after
        assert "JOIN langchain_pg_collection" not in sql
        assert "e.collection_id = (SELECT uuid FROM langchain_pg_collection" in sql
        assert conn.execute.call_args.args[1]["collection"] == "rules"
        assert conn.execute.call_args.args[1]["codes"] == ["G0556", "G0557"]
        assert conn.execute.call_args.args[1]["window"] == 4
        assert [s["snippet"] for s in out] == [
            "para one G0556",
            "para two G0556",
            "para four G0557",
        ]
        assert all(s["kind"] == "rule" and s["score"] == 0.0 for s in out)

    def test_snippets_are_capped(self):
        engine, _ = _engine([("x" * 900, _rule_md(1, "G0556"))])
        (out,) = code_cited_sources(engine, ["G0556"], per_code=2)
        assert len(out["snippet"]) == 250

    def test_engine_error_yields_empty(self, caplog):
        engine = MagicMock()
        engine.begin.side_effect = RuntimeError("pg down")
        assert code_cited_sources(engine, ["G0556"], per_code=2) == []
        assert "code-cited sources skipped" in caplog.text

    def test_no_codes_returns_empty_without_querying(self):
        engine = MagicMock()
        assert code_cited_sources(engine, [], per_code=2) == []
        engine.begin.assert_not_called()


class TestCodeCitedSourcesMultiCollection:
    def test_concatenated_rules_then_comments_then_corpus(self):
        engine, _conn = _multi_engine(
            by_code={
                "rules": [("rule chunk", _rule_md(1, "G0556"))],
                "comments": [("comment chunk", _comment_md(1, "G0556"))],
                "corpus": [("corpus chunk", _corpus_md(1, "G0556"))],
            }
        )
        out = code_cited_sources(
            engine, ["G0556"], per_code=2, collections=("rules", "comments", "corpus")
        )
        assert [s["kind"] for s in out] == ["rule", "comment", "corpus"]

    def test_per_collection_cap_applies_independently(self):
        engine, _conn = _multi_engine(
            by_code={
                "rules": [
                    ("r1", _rule_md(1, "G0556")),
                    ("r2", _rule_md(2, "G0556")),
                    ("r3", _rule_md(3, "G0556")),  # over the per-collection cap
                ],
                "comments": [
                    ("c1", _comment_md(1, "G0556")),
                    ("c2", _comment_md(2, "G0556")),
                ],
            }
        )
        out = code_cited_sources(
            engine, ["G0556"], per_code=2, collections=("rules", "comments")
        )
        # Capped to 2/collection, then round-robin interleaved across
        # collections (rules, comments), not simply concatenated.
        assert [s["snippet"] for s in out] == ["r1", "c1", "r2", "c2"]

    def test_dedupe_uses_seq_for_non_rule_kinds(self):
        engine, _conn = _multi_engine(
            by_code={
                "comments": [
                    ("c1", _comment_md(1, "G0556", item_key="C1")),
                    ("c1 dup", _comment_md(1, "G0556", item_key="C1")),  # same seq
                    ("c2", _comment_md(2, "G0556", item_key="C1")),  # different seq
                ]
            }
        )
        out = code_cited_sources(
            engine, ["G0556"], per_code=5, collections=("comments",)
        )
        assert [s["snippet"] for s in out] == ["c1", "c2"]


class TestFamilyCitedSources:
    def test_skipped_entirely_when_no_families(self):
        engine, conn = _multi_engine(by_code={"rules": [("r1", _rule_md(1, "G0556"))]})
        code_cited_sources(engine, ["G0556"], per_code=2, collections=("rules",))
        assert conn.execute.call_count == 1  # codes query only, no family round trip

    def test_sql_shape_and_params(self):
        engine, conn = _multi_engine(by_family={"rules": []})
        code_cited_sources(
            engine, [], per_code=2, collections=("rules",), families=["CCM"]
        )
        sql = str(conn.execute.call_args.args[0])
        assert "&&" in sql
        # one chunk per (family, item_key), then the survivors ranked per family
        assert sql.count("PARTITION BY") == 2
        assert "PARTITION BY w.family, e.cmetadata->>'item_key'" in sql
        assert "PARTITION BY family" in sql
        assert conn.execute.call_args.args[1] == {
            "collection": "rules",
            "families": ["CCM"],
            "window": 4,
        }

    def test_called_once_per_collection_with_same_params(self):
        engine, conn = _multi_engine(
            by_family={"rules": [], "comments": [], "corpus": []}
        )
        code_cited_sources(
            engine,
            [],
            per_code=3,
            collections=("rules", "comments", "corpus"),
            families=["CCM"],
        )
        assert conn.execute.call_count == 3
        for call, collection in zip(
            conn.execute.call_args_list, ("rules", "comments", "corpus")
        ):
            assert call.args[1] == {
                "collection": collection,
                "families": ["CCM"],
                "window": 6,
            }

    def test_family_rows_appended_after_code_rows(self):
        engine, _conn = _multi_engine(
            by_code={"rules": [("code hit", _rule_md(1, "G0556"))]},
            by_family={
                "corpus": [
                    (
                        "CCM guideline",
                        _corpus_md(1, codes="", families="CCM", item_key="CPT-ED"),
                    )
                ]
            },
        )
        out = code_cited_sources(
            engine,
            ["G0556"],
            per_code=2,
            collections=("rules", "corpus"),
            families=["CCM"],
        )
        assert [s["snippet"] for s in out] == ["code hit", "CCM guideline"]
        assert out[1]["kind"] == "corpus"

    def test_family_dedupes_against_a_code_cited_chunk(self):
        md = _rule_md(1, "G0556")
        md["families"] = "APCM"
        engine, _conn = _multi_engine(
            by_code={"rules": [("shared", md)]}, by_family={"rules": [("shared", md)]}
        )
        out = code_cited_sources(
            engine, ["G0556"], per_code=2, collections=("rules",), families=["APCM"]
        )
        assert [s["snippet"] for s in out] == ["shared"]


class TestMaxTotal:
    """``max_total`` interleaves code/family rows round-robin across
    collections, then truncates (#budget)."""

    def test_interleaves_three_collections_round_robin(self):
        engine, _conn = _multi_engine(
            by_code={
                "rules": [("r1", _rule_md(1, "G0556")), ("r2", _rule_md(2, "G0556"))],
                "comments": [
                    ("c1", _comment_md(1, "G0556")),
                    ("c2", _comment_md(2, "G0556")),
                ],
                "corpus": [
                    ("x1", _corpus_md(1, codes="G0556", item_key="X1")),
                    ("x2", _corpus_md(2, codes="G0556", item_key="X2")),
                ],
            }
        )
        out = code_cited_sources(
            engine,
            ["G0556"],
            per_code=5,
            collections=("rules", "comments", "corpus"),
            max_total=0,
        )
        assert [s["snippet"] for s in out] == ["r1", "c1", "x1", "r2", "c2", "x2"]

    def test_family_rows_interleaved_then_appended_after_code_rows(self):
        engine, _conn = _multi_engine(
            by_code={
                "rules": [("r1", _rule_md(1, "G0556"))],
                "comments": [("c1", _comment_md(1, "G0556"))],
            },
            by_family={
                "rules": [
                    (
                        "f-rules",
                        _corpus_md(9, codes="", families="CCM", item_key="FR"),
                    )
                ],
                "comments": [
                    (
                        "f-comments",
                        _corpus_md(10, codes="", families="CCM", item_key="FC"),
                    )
                ],
            },
        )
        out = code_cited_sources(
            engine,
            ["G0556"],
            per_code=5,
            collections=("rules", "comments"),
            families=["CCM"],
            max_total=0,
        )
        assert [s["snippet"] for s in out] == ["r1", "c1", "f-rules", "f-comments"]

    def test_truncates_to_max_total(self):
        engine, _conn = _multi_engine(
            by_code={
                "rules": [("r1", _rule_md(1, "G0556")), ("r2", _rule_md(2, "G0556"))],
                "comments": [
                    ("c1", _comment_md(1, "G0556")),
                    ("c2", _comment_md(2, "G0556")),
                ],
            }
        )
        out = code_cited_sources(
            engine,
            ["G0556"],
            per_code=5,
            collections=("rules", "comments"),
            max_total=3,
        )
        assert [s["snippet"] for s in out] == ["r1", "c1", "r2"]

    def test_max_total_zero_or_negative_is_unlimited(self):
        rows = [(f"r{i}", _rule_md(i, "G0556")) for i in range(1, 15)]
        engine, _conn = _multi_engine(by_code={"rules": rows})
        for unlimited in (0, -1):
            out = code_cited_sources(
                engine,
                ["G0556"],
                per_code=20,
                collections=("rules",),
                max_total=unlimited,
            )
            assert len(out) == 14

    def test_default_max_total_is_twelve(self):
        rows = [(f"r{i}", _rule_md(i, "G0556")) for i in range(1, 15)]
        engine, _conn = _multi_engine(by_code={"rules": rows})
        out = code_cited_sources(engine, ["G0556"], per_code=20, collections=("rules",))
        assert len(out) == 12


def _comment_docket_md(
    seq, codes, docket, item_key=None, date="2025-01-01", year="2025"
):
    return {
        "kind": "comment",
        "item_key": item_key or f"{docket}-C{seq}",
        "seq": str(seq),
        "comment_id": f"{docket}-{seq}",
        "date": date,
        "year": year,
        "docket": docket,
        "codes": codes,
        "families": "",
    }


def _docket_fake_engine(docket_rows, code_rows_by_collection=None):
    """A fake engine that tells the docket-SQL call apart from the
    normal per-code window call by the params shape alone — the
    per-code window (``_collect``) always sends a ``window`` bind
    param; the docket window (``_collect_by_docket``) never does."""
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    code_rows_by_collection = code_rows_by_collection or {}

    def _execute(_sql, params):
        result = MagicMock()
        if "window" not in params:
            result.fetchall.return_value = docket_rows
        else:
            result.fetchall.return_value = code_rows_by_collection.get(
                params["collection"], []
            )
        return result

    conn.execute.side_effect = _execute
    return engine, conn


class TestCodeCitedSourcesByDocket:
    """Ruling B7 — ``by_docket`` swaps the comments collection's
    per-code recency window for one chunk per docket, all dockets, in
    timeline mode only."""

    def test_docket_sql_used_only_for_comments(self):
        engine, conn = _docket_fake_engine(
            docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))],
            code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]},
        )
        code_cited_sources(
            engine,
            ["G2211"],
            per_code=3,
            collections=("rules", "comments"),
            by_docket=True,
        )
        calls = conn.execute.call_args_list
        comments_calls = [c for c in calls if c.args[1]["collection"] == "comments"]
        rules_calls = [c for c in calls if c.args[1]["collection"] == "rules"]
        assert comments_calls and "window" not in comments_calls[0].args[1]
        assert rules_calls and "window" in rules_calls[0].args[1]

    def test_by_docket_false_uses_the_normal_window_for_comments(self):
        engine, conn = _docket_fake_engine(
            docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))],
            code_rows_by_collection={"comments": [("c1", _comment_md(1, "G2211"))]},
        )
        code_cited_sources(
            engine,
            ["G2211"],
            per_code=3,
            collections=("comments",),
            by_docket=False,
        )
        calls = conn.execute.call_args_list
        assert calls and "window" in calls[0].args[1]

    def test_by_docket_without_comments_in_collections_is_a_noop(self):
        engine, conn = _docket_fake_engine(
            docket_rows=[("d1", _comment_docket_md(1, "G2211", "CMS-2023-0121"))],
            code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]},
        )
        code_cited_sources(
            engine, ["G2211"], per_code=3, collections=("rules",), by_docket=True
        )
        calls = conn.execute.call_args_list
        assert calls and "window" in calls[0].args[1]

    def test_three_dockets_survive_dedupe_and_are_interleaved(self):
        docket_rows = [
            ("c23", _comment_docket_md(1, "G2211", "CMS-2023-0121", item_key="C23")),
            ("c25", _comment_docket_md(1, "G2211", "CMS-2025-0304", item_key="C25")),
            ("c26", _comment_docket_md(1, "G2211", "CMS-2026-2377", item_key="C26")),
        ]
        engine, _conn = _docket_fake_engine(
            docket_rows=docket_rows,
            code_rows_by_collection={"rules": [("r1", _rule_md(1, "G2211"))]},
        )
        out = code_cited_sources(
            engine,
            ["G2211"],
            per_code=3,
            collections=("rules", "comments"),
            by_docket=True,
        )
        dockets = {s["docket"] for s in out if s["kind"] == "comment"}
        assert dockets == {"CMS-2023-0121", "CMS-2025-0304", "CMS-2026-2377"}
        # round-robin interleaved with the rules row, not appended en bloc
        assert [s["kind"] for s in out] == ["rule", "comment", "comment", "comment"]


class TestCollectByDocket:
    """``_collect_by_docket`` directly — the per-docket window behind
    ``by_docket=True`` (Ruling B7)."""

    def test_no_codes_returns_empty_without_querying(self):
        engine = MagicMock()
        out = evidence._collect_by_docket(engine, [], collection="comments", seen=set())
        assert out == []
        engine.begin.assert_not_called()

    def test_engine_error_yields_empty(self, caplog):
        engine = MagicMock()
        engine.begin.side_effect = RuntimeError("pg down")
        out = evidence._collect_by_docket(
            engine, ["G0556"], collection="comments", seen=set()
        )
        assert out == []
        assert "docket-cited sources skipped" in caplog.text

    def test_dedupe_against_seen_skips_the_row(self):
        md = _comment_docket_md(1, "G0556", "CMS-2023-0121", item_key="C1")
        engine, _conn = _engine([("c1", md)])
        seen = {evidence._dedupe_key(md)}
        out = evidence._collect_by_docket(
            engine, ["G0556"], collection="comments", seen=seen
        )
        assert out == []

    def test_caps_at_docket_max(self):
        rows = [
            (
                f"c{i}",
                _comment_docket_md(1, "G0556", f"CMS-2023-{i:04d}", item_key=f"C{i}"),
            )
            for i in range(evidence._DOCKET_MAX + 8)
        ]
        engine, _conn = _engine(rows)
        out = evidence._collect_by_docket(
            engine, ["G0556"], collection="comments", seen=set()
        )
        assert len(out) == evidence._DOCKET_MAX


class TestMergeSources:
    def test_keeps_order_and_dedupes_by_label(self):
        a = [{"label": "X", "score": 0.9}, {"label": "Y", "score": 0.8}]
        b = [{"label": "Y", "score": 0.0}, {"label": "Z", "score": 0.0}]
        assert [s["label"] for s in merge_sources(a, b)] == ["X", "Y", "Z"]
        assert merge_sources(a, b)[1]["score"] == 0.8

    def test_rule_rows_dedupe_by_item_key_p_id_even_with_different_labels(self):
        # Ruling B13 (I2): two rule chunks for the same paragraph can now
        # carry different labels (a retrieved "85 FR 84639 ¶12" vs a
        # lineage "CY2021 PFS final 85 FR 84639 ¶1578") — the pair must
        # still dedupe on (item_key, p_id), not slip through on label.
        a = [
            {
                "label": "85 FR 84639 ¶12",
                "kind": "rule",
                "item_key": "YBM4IZUS",
                "p_id": "1578",
            }
        ]
        b = [
            {
                "label": "CY2021 PFS final 85 FR 84639 ¶1578",
                "kind": "rule",
                "item_key": "YBM4IZUS",
                "p_id": "1578",
            }
        ]
        out = merge_sources(a, b)
        assert len(out) == 1
        assert out[0]["label"] == "85 FR 84639 ¶12"

    def test_rule_row_missing_item_key_or_p_id_falls_back_to_label(self):
        a = [{"label": "some rule", "kind": "rule", "item_key": "", "p_id": ""}]
        b = [{"label": "some rule", "kind": "rule", "item_key": "K", "p_id": ""}]
        c = [{"label": "other", "kind": "rule", "item_key": "K", "p_id": ""}]
        out = merge_sources(merge_sources(a, b), c)
        assert [s["label"] for s in out] == ["some rule", "other"]

    def test_non_rule_kinds_still_dedupe_by_label(self):
        a = [{"label": "L", "kind": "comment", "item_key": "K", "p_id": ""}]
        b = [{"label": "L", "kind": "comment", "item_key": "K", "p_id": ""}]
        assert len(merge_sources(a, b)) == 1

    def test_extra_rule_row_with_a_new_key_is_kept(self):
        a = [{"label": "A", "kind": "rule", "item_key": "K1", "p_id": "1"}]
        b = [{"label": "B", "kind": "rule", "item_key": "K2", "p_id": "9"}]
        out = merge_sources(a, b)
        assert [s["label"] for s in out] == ["A", "B"]


def _cpt_section_row(con, year, item_key, sec_id, title, path_key, guideline):
    con.execute(
        "INSERT INTO pfs.cpt_section VALUES (?,?,?,?,?,?,?,?,?,?)",
        [
            year,
            item_key,
            sec_id,
            3,
            title,
            path_key.split(" > "),
            path_key,
            "",
            "",
            guideline,
        ],
    )


def _family_row(con, key, note, code="99490"):
    con.execute(
        "INSERT INTO pfs.code_family VALUES (?,?,?,?,?,?,?,?,?)",
        [key, key, code, "member", None, None, "", 0, note],
    )


CCM_PATH = "Evaluation and Management > Care Management Services > Chronic Care Management Services"
PARENT_PATH = "Surgery > General"
LEAF_PATH = "Surgery > General > Leaf With No Guideline"


class TestManualSources:
    """``manual_sources`` — the CPT manual's own guideline text for a
    detected family's heading, straight from the DuckDB replica (#691
    Ruling B4: never a pgvector query — the CPT chunks' ``section``
    metadata is junk EPUB headings, not the family heading)."""

    @pytest.fixture
    def cpt_replica(self):
        con = duckdb.connect(":memory:")
        ensure_tables(con)
        # CCM: two editions have the heading's guideline text; the newer
        # (2024) one must win over the older (2022) one.
        _family_row(con, "CCM", CCM_PATH)
        _cpt_section_row(
            con,
            2022,
            "OLDED001",
            "S0",
            "Chronic Care Management Services",
            CCM_PATH,
            "2022 text: non-face-to-face care management services.",
        )
        _cpt_section_row(
            con,
            2024,
            "GQGTPGYV",
            "S1",
            "Chronic Care Management Services",
            CCM_PATH,
            "Chronic care management services are non-face-to-face "
            "services provided to a patient with two or more chronic conditions.",
        )
        # PARENTONLY: the leaf heading has no guideline of its own, but
        # its one-level parent does — the fallback must pick it up.
        _family_row(con, "PARENTONLY", LEAF_PATH)
        _cpt_section_row(
            con, 2024, "LEAFED001", "S2", "Leaf With No Guideline", LEAF_PATH, ""
        )
        _cpt_section_row(
            con,
            2023,
            "PARED0001",
            "S3",
            "General",
            PARENT_PATH,
            "General guidance text for the parent heading.",
        )
        # NONOTE: a real family row, but its note is empty.
        _family_row(con, "NONOTE", "", code="99999")
        yield con
        con.close()

    @pytest.fixture(autouse=True)
    def _fake_store(self):
        items = {
            "GQGTPGYV": SimpleNamespace(
                date_published="2023-11-16",
                url="https://ebooks.ama-assn.org/cpt2024",
            )
        }

        def _get(key):
            if key in items:
                return items[key]
            raise KeyError(key)

        store = MagicMock()
        store.get.side_effect = _get
        with patch("llm.evidence._store", return_value=store):
            yield store

    def test_newest_edition_with_guideline_wins(self, cpt_replica):
        out = manual_sources(cpt_replica, ["CCM"])
        assert len(out) == 1
        s = out[0]
        assert s["kind"] == "corpus"
        assert s["title"] == "Chronic Care Management Services — CPT 2024"
        assert s["item_key"] == "GQGTPGYV" and s["seq"] == "S1"
        assert s["section"] == "Chronic Care Management Services"
        assert s["date"] == "2023-11-16"
        assert s["url"] == "https://ebooks.ama-assn.org/cpt2024"
        assert s["snippet"].startswith("Chronic care management services")

    def test_parent_fallback_when_leaf_has_no_guideline(self, cpt_replica):
        out = manual_sources(cpt_replica, ["PARENTONLY"])
        assert len(out) == 1
        s = out[0]
        assert s["title"] == "General — CPT 2023"
        assert s["item_key"] == "PARED0001" and s["seq"] == "S3"
        assert "General guidance text" in s["snippet"]
        # unresolved bib item (not in the fake store) — falls back to
        # the plain edition year and an empty url.
        assert s["date"] == "2023-01-01" and s["url"] == ""

    def test_family_with_empty_note_yields_nothing(self, cpt_replica):
        assert manual_sources(cpt_replica, ["NONOTE"]) == []

    def test_family_with_no_family_row_yields_nothing(self, cpt_replica):
        assert manual_sources(cpt_replica, ["GHOST"]) == []

    def test_no_families_yields_nothing_without_querying(self, cpt_replica):
        assert manual_sources(cpt_replica, []) == []

    def test_per_family_cap(self, cpt_replica):
        complex_path = (
            "Evaluation and Management > Care Management Services > "
            "Complex Chronic Care Management Services"
        )
        _family_row(cpt_replica, "CCM", complex_path, code="99487")
        _cpt_section_row(
            cpt_replica,
            2024,
            "GQGTPGYV",
            "S4",
            "Complex Chronic Care Management Services",
            complex_path,
            "Complex chronic care management guideline text.",
        )
        assert len(manual_sources(cpt_replica, ["CCM"], per_family=1)) == 1
        assert len(manual_sources(cpt_replica, ["CCM"], per_family=2)) == 2

    def test_total_cap_across_families(self, cpt_replica):
        # Ruling B11: even with three families each good for a source,
        # manual_sources never returns more than max_total (default 2).
        _family_row(cpt_replica, "PARENTONLY2", PARENT_PATH, code="99498")
        out = manual_sources(cpt_replica, ["CCM", "PARENTONLY", "PARENTONLY2"])
        assert len(out) == 2
        assert len(manual_sources(cpt_replica, ["CCM", "PARENTONLY"], max_total=1)) == 1

    def test_snippet_chars_bounds_the_text_passed_to_as_source(self, cpt_replica):
        out = manual_sources(cpt_replica, ["CCM"], snippet_chars=10)
        assert len(out[0]["snippet"]) <= 10

    def test_missing_tables_yield_empty(self):
        con = duckdb.connect(":memory:")
        try:
            assert manual_sources(con, ["CCM"]) == []
        finally:
            con.close()

    def test_guideline_row_returns_none_when_no_heading_matches(self, cpt_replica):
        # Neither the leaf path nor its one-level parent exists in
        # pfs.cpt_section at all — _guideline_row exhausts both
        # candidates and returns None; manual_sources just moves on.
        _family_row(cpt_replica, "NOMATCH", "Foo > Bar > Baz", code="00001")
        assert manual_sources(cpt_replica, ["NOMATCH"]) == []

    def test_family_notes_non_missing_table_error_is_logged(self, cpt_replica, caplog):
        with patch("llm.evidence._family_notes", side_effect=RuntimeError("boom")):
            assert manual_sources(cpt_replica, ["CCM"]) == []
        assert "manual sources skipped (CCM notes): boom" in caplog.text

    def test_guideline_row_non_missing_table_error_is_logged(self, cpt_replica, caplog):
        with patch("llm.evidence._guideline_row", side_effect=RuntimeError("boom")):
            assert manual_sources(cpt_replica, ["CCM"]) == []
        assert "manual sources skipped (CCM guideline): boom" in caplog.text