stack/tests/llm/test_lineage.py

"""llm.lineage — dated timeline, element diffs and guidance as a
lineage SSE event + cited prompt block."""

from __future__ import annotations

import json
import sqlite3
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from types import SimpleNamespace
from unittest.mock import patch

import duckdb
import pytest

import llm.lineage as lineage
from llm.config import LlmConfig
from llm.lineage import (
    ElementDiff,
    GuidanceRef,
    LineageEvent,
    LineageEvidence,
    _reconcile_labels,
    lineage_evidence,
    lineage_sources,
    rule_label,
)
from pfs.codetables import (
    ElementRow,
    EventRow,
    GuidanceRow,
    ensure_tables,
    write_elements,
    write_events,
    write_guidance,
)
from pfs.families import Detection

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",
    lineage_max_rows=25,
    lineage_on_demand_max=3,
)


@pytest.fixture(autouse=True)
def _clear_caches():
    """Every process-global cache lineage.py keeps is shared across
    tests (and with llm.evidence's replica cache) — no test may inherit
    another's handle, store, or on-demand memo."""
    import llm.evidence as evidence

    def _reset():
        evidence._REPLICA = None
        lineage._STORE_LOCAL = threading.local()
        lineage._ITEM_CACHE.reset()
        lineage._URL_CACHE.reset()
        lineage._ON_DEMAND_CACHE.reset()
        lineage._PARAGRAPH_CACHE.reset()

    _reset()
    yield
    _reset()


class _FakeStore:
    """A bib.Store stand-in: ``items``/``fr_anchors``/``fr_anchor_docs``
    tables the way bib.frlink.resolve and Store.get expect, seeded with
    the real CY2015/CY2021 CCM history (99490 adopted in 2015, G2058
    replaced by the new CPT code 99439 in 2021) already used as a fixture
    in tests/pfs/test_lineage.py."""

    def __init__(self):
        self.con = sqlite3.connect(":memory:")
        self.con.row_factory = sqlite3.Row
        self.con.executescript(
            "CREATE TABLE items (key TEXT PRIMARY KEY, title TEXT, "
            "date_published TEXT, url TEXT);"
            "CREATE TABLE fr_anchors (item_key TEXT, p_id INTEGER, "
            "page INTEGER, ordinal INTEGER, text TEXT);"
            "CREATE TABLE fr_anchor_docs (item_key TEXT, html_url TEXT, "
            "start_page INTEGER, end_page INTEGER, fr_volume INTEGER);"
        )
        self.con.executemany(
            "INSERT INTO items VALUES (?,?,?,?)",
            [
                (
                    "DE2VH9PD",
                    "Medicare Program; CY 2015 PFS Final Rule",
                    "2014-11-13",
                    "https://www.federalregister.gov/d/2014-2015doc",
                ),
                (
                    "YBM4IZUS",
                    "Medicare Program; CY 2021 Payment Policies Under the PFS",
                    "2020-12-28",
                    "https://www.federalregister.gov/d/2020-2021doc",
                ),
                (
                    "ZPROP2027",
                    "Medicare Program; CY 2027 Proposed Payment Policies",
                    "2026-07-16",
                    "https://www.federalregister.gov/d/2026-2027doc",
                ),
            ],
        )
        self.con.executemany(
            "INSERT INTO fr_anchor_docs VALUES (?,?,?,?,?)",
            [
                ("DE2VH9PD", "https://fr.test/2015-doc", 67000, 68000, 79),
                ("YBM4IZUS", "https://fr.test/2021-doc", 84000, 85000, 85),
            ],
        )
        self.con.executemany(
            "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
            [
                (
                    "DE2VH9PD",
                    1251,
                    67716,
                    1,
                    "Accordingly, we will adopt CPT code 99490 for Medicare "
                    "CCM services.",
                ),
                (
                    "YBM4IZUS",
                    686,
                    84547,
                    1,
                    "We are finalizing HCPCS code G2058 as new CPT code 99439.",
                ),
                (
                    "YBM4IZUS",
                    1578,
                    84639,
                    2,
                    "A temporary crosswalk between G2058 and new CPT code "
                    "99439 (with a descriptor identical to G2058).",
                ),
                # The proposed rule decides its kind from its own text —
                # bib.frlink.rule_kind counts "this proposed rule"/"we
                # propose" against "this final rule"/"we are finalizing".
                (
                    "ZPROP2027",
                    394,
                    32388,
                    1,
                    "In this proposed rule we propose to keep the three-step "
                    "telehealth review.",
                ),
            ],
        )

    def _con(self):
        return self.con

    def get(self, key: str) -> SimpleNamespace:
        row = self.con.execute(
            "SELECT title, date_published FROM items WHERE key = ?", (key,)
        ).fetchone()
        if row is None:
            raise KeyError(key)
        return SimpleNamespace(title=row["title"], date_published=row["date_published"])

    def close(self):
        self.con.close()


@pytest.fixture
def store():
    s = _FakeStore()
    yield s
    s.close()


@pytest.fixture
def con(tmp_path):
    """A real temp DuckDB replica (not :memory:) so lineage_evidence's
    real _connect/_replica_path path is exercised end to end, the same
    way TestConcurrentChats in test_evidence.py does."""
    db = tmp_path / "aco.ro.duckdb"
    c = duckdb.connect(str(db))
    ensure_tables(c)
    yield c, db
    c.close()


def _ev(
    code,
    year,
    kind,
    frm="",
    to="",
    item_key="",
    p_id=0,
    page=0,
    source="fr",
    anchored=True,
    note="",
) -> EventRow:
    return EventRow(
        code, year, kind, frm, to, item_key, p_id, page, source, anchored, note
    )


def _el(code, year, type_, value, item_key="", p_id=0, source="fr") -> ElementRow:
    return ElementRow(code, year, type_, value, "", "", item_key, p_id, 0, source)


def _gd(
    family, code, kind, locator, item_key_src, p_id_src, item_key=""
) -> GuidanceRow:
    return GuidanceRow(family, code, kind, locator, item_key, item_key_src, p_id_src, 0)


class TestKindRank:
    def test_unknown_kind_sorts_after_every_known_kind(self):
        """``_kind_rank`` is only ever called on ``EventRow.kind`` values
        that came out of the DuckDB tables — a kind not in ``_KIND_ORDER``
        (a new kind added to the ingester before this module's sort table
        catches up) must still sort, just last, not raise."""
        assert lineage._kind_rank("some_future_kind") == len(lineage._KIND_ORDER)


class TestRuleLabel:
    def test_title_without_a_kind_word_is_rule_not_final(self):
        """Since CY2018 the CMS titles say neither "Proposed" nor
        "Final"; without a decided kind the label must not claim
        ``final`` (a proposed rule presented as final is a false
        statement of policy)."""
        assert (
            rule_label(
                "Medicare Program; CY 2021 Payment Policies Under the PFS",
                "2020-12-28",
                1578,
            )
            == "CY2021 PFS rule ¶1578"
        )

    def test_final_from_title(self):
        assert (
            rule_label("Medicare Program; CY 2015 PFS Final Rule", "2014-11-13", 1251)
            == "CY2015 PFS final ¶1251"
        )

    def test_kind_from_the_rule_text_overrides_the_title(self):
        """The chat passes ``kind`` decided by ``bib.frlink.rule_kind``
        from the rule's own paragraphs — a title that says nothing
        still labels a proposed rule ``proposed``."""
        assert (
            rule_label(
                "Medicare and Medicaid Programs; CY 2027 Payment Policies",
                "2026-07-16",
                394,
                kind="proposed",
            )
            == "CY2027 PFS proposed ¶394"
        )

    def test_correction_from_title(self):
        assert (
            rule_label("Medicare Program; CY 2016 PFS Correction", "2016-03-08", 12)
            == "CY2016 PFS correction ¶12"
        )

    def test_proposed(self):
        assert (
            rule_label(
                "Medicare Program; CY 2027 Proposed Payment Policies",
                "2026-07-16",
                394,
            )
            == "CY2027 PFS proposed ¶394"
        )

    def test_final_with_volume_and_page(self):
        assert (
            rule_label(
                "Medicare Program; CY 2021 Payment Policies Under the PFS",
                "2020-12-28",
                1578,
                vol="85",
                page=84639,
                kind="final",
            )
            == "CY2021 PFS final 85 FR 84639 ¶1578"
        )

    def test_missing_page_falls_back_to_the_short_form(self):
        assert (
            rule_label(
                "Medicare Program; CY 2021 Payment Policies Under the PFS",
                "2020-12-28",
                1578,
                vol="85",
                kind="final",
            )
            == "CY2021 PFS final ¶1578"
        )


class TestCollapse:
    def test_prefers_anchored_fr_over_rvu_for_the_same_group(self):
        rows = [
            _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686, source="fr"),
            _ev("99439", 2021, "created", source="rvu", anchored=True, note="dupe"),
        ]
        collapsed = lineage._collapse(rows)
        assert len(collapsed) == 1
        assert collapsed[0].source == "fr" and collapsed[0].p_id == 686

    def test_tiebreaks_on_lowest_p_id(self):
        rows = [
            _ev(
                "99439",
                2021,
                "replaces",
                to="",
                frm="G2058",
                item_key="YBM4IZUS",
                p_id=2999,
                source="fr",
            ),
            _ev(
                "99439",
                2021,
                "replaces",
                frm="G2058",
                item_key="YBM4IZUS",
                p_id=1578,
                source="fr",
            ),
        ]
        collapsed = lineage._collapse(rows)
        assert len(collapsed) == 1 and collapsed[0].p_id == 1578

    def test_sorted_by_year_code_then_kind_order(self):
        rows = [
            _ev("99439", 2021, "replaces", frm="G2058", p_id=1),
            _ev("G2058", 2021, "replaced_by", to="99439", p_id=1),
            _ev("99439", 2015, "created", p_id=2),
        ]
        collapsed = lineage._collapse(rows)
        assert [(e.year, e.code, e.kind) for e in collapsed] == [
            (2015, "99439", "created"),
            (2021, "99439", "replaces"),
            (2021, "G2058", "replaced_by"),
        ]

    def test_distinct_groups_are_not_merged(self):
        rows = [
            _ev("99439", 2021, "created", p_id=1),
            _ev("99439", 2021, "replaces", frm="G2058", p_id=1),
        ]
        assert len(lineage._collapse(rows)) == 2

    def test_created_across_three_years_collapses_to_the_earliest(self):
        """Ruling B5: G2058 is mentioned as "created" by three different
        rules (2021, 2022, 2023 — each rule re-establishing the G-code
        after it was crosswalked away) but it was only actually created
        once."""
        rows = [
            _ev("G2058", 2021, "created", item_key="A", p_id=100),
            _ev("G2058", 2022, "created", item_key="B", p_id=200),
            _ev("G2058", 2023, "created", item_key="C", p_id=300),
        ]
        collapsed = lineage._collapse(rows)
        assert [(e.year, e.code, e.kind) for e in collapsed] == [
            (2021, "G2058", "created")
        ]
        assert collapsed[0].item_key == "A"

    def test_deleted_twice_keeps_only_the_first(self):
        rows = [
            _ev("99213", 2019, "deleted", item_key="A", p_id=1),
            _ev("99213", 2024, "deleted", item_key="B", p_id=2),
        ]
        collapsed = lineage._collapse(rows)
        assert [(e.year, e.kind) for e in collapsed] == [(2019, "deleted")]

    def test_earliest_only_picks_the_preferred_row_within_that_year(self):
        """Two rows tie on the earliest year — the anchored-FR/lowest-p_id
        preference still applies among them."""
        rows = [
            _ev(
                "G2058",
                2021,
                "created",
                item_key="LATE",
                p_id=999,
                source="fr",
                anchored=True,
            ),
            _ev(
                "G2058",
                2021,
                "created",
                item_key="EARLY",
                p_id=100,
                source="fr",
                anchored=True,
            ),
            _ev("G2058", 2022, "created", item_key="OTHER", p_id=1),
        ]
        collapsed = lineage._collapse(rows)
        assert len(collapsed) == 1
        assert collapsed[0].year == 2021 and collapsed[0].item_key == "EARLY"

    def test_replaced_by_rows_in_two_years_both_survive(self):
        """replaced_by is directional/repeatable, not a point-in-time
        kind — a code can legitimately be replaced by different things
        in different years, so both rows must survive collapse."""
        rows = [
            _ev("99439", 2021, "replaced_by", to="G2058", p_id=1),
            _ev("99439", 2022, "replaced_by", to="G2214", p_id=2),
        ]
        collapsed = lineage._collapse(rows)
        assert [(e.year, e.to_codes) for e in collapsed] == [
            (2021, "G2058"),
            (2022, "G2214"),
        ]

    def test_final_beats_proposed_on_a_tie(self, store):
        """Ruling B12: two rows tied on (code, year, kind, from, to,
        p_id) — one from a proposed rule (ZPROP2027), one from a final
        rule (YBM4IZUS) — the final rule wins."""
        rows = [
            _ev(
                "99439",
                2021,
                "created",
                item_key="ZPROP2027",
                p_id=999,
                source="fr",
                anchored=True,
            ),
            _ev(
                "99439",
                2021,
                "created",
                item_key="YBM4IZUS",
                p_id=999,
                source="fr",
                anchored=True,
            ),
        ]
        collapsed = lineage._collapse(rows, store, 0)
        assert len(collapsed) == 1
        assert collapsed[0].item_key == "YBM4IZUS"

    def test_fully_tied_pair_decided_by_item_key(self, store):
        """Neither row resolves to a proposed title (both unknown item
        keys) — the tie is broken deterministically by item_key."""
        rows = [
            _ev(
                "99439",
                2021,
                "created",
                item_key="ZZZZ0000",
                p_id=999,
                source="fr",
                anchored=True,
            ),
            _ev(
                "99439",
                2021,
                "created",
                item_key="AAAA0000",
                p_id=999,
                source="fr",
                anchored=True,
            ),
        ]
        collapsed = lineage._collapse(rows, store, 0)
        assert len(collapsed) == 1
        assert collapsed[0].item_key == "AAAA0000"

    def test_no_store_never_prefers_proposed_but_stays_deterministic(self):
        """Without a store (the pre-B12 call shape), every row is
        treated as not-proposed — the anchored/p_id/item_key ordering
        still applies and never raises."""
        rows = [
            _ev("99439", 2021, "created", item_key="B", p_id=999, source="fr"),
            _ev("99439", 2021, "created", item_key="A", p_id=999, source="fr"),
        ]
        collapsed = lineage._collapse(rows)
        assert len(collapsed) == 1
        assert collapsed[0].item_key == "A"


class TestElementDiffs:
    def test_no_diffs_when_fewer_than_two_codes_have_elements(self, con):
        c, _ = con
        assert lineage._element_diffs(c, ["99490"], None, 1) == ((), "")

    def test_notes_when_exactly_one_code_has_elements(self, con):
        c, _ = con
        write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
        diffs, note = lineage._element_diffs(c, ["99490", "99491", "G0556"], None, 1)
        assert diffs == ()
        assert note == "elements extracted for 1 of 3 codes"

    def test_shared_value_diffs_from_the_one_missing_code(self, con, store):
        c, _ = con
        write_elements(
            c,
            "99490",
            [_el("99490", 2026, "consent", "required", "YBM4IZUS", 1578)],
            [],
        )
        write_elements(c, "99491", [_el("99491", 2026, "consent", "required")], [])
        write_elements(c, "G0556", [_el("G0556", 2026, "telehealth", "yes")], [])
        diffs, note = lineage._element_diffs(c, ["99490", "99491", "G0556"], store, 1)
        assert note == ""
        by_pair = {(d.type, d.value): d for d in diffs}
        assert set(by_pair) == {("consent", "required"), ("telehealth", "yes")}
        consent = by_pair[("consent", "required")]
        assert consent.in_codes == ("99490", "99491")
        assert consent.not_in_codes == ("G0556",)
        assert (
            consent.label == "CY2021 PFS final 85 FR 84639 ¶1578"
        )  # from the anchored 99490 row
        telehealth = by_pair[("telehealth", "yes")]
        assert telehealth.in_codes == ("G0556",)
        assert telehealth.not_in_codes == ("99490", "99491")

    def test_no_diff_when_every_code_with_elements_shares_the_value(self, con):
        c, _ = con
        write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
        write_elements(c, "99491", [_el("99491", 2026, "consent", "required")], [])
        diffs, _ = lineage._element_diffs(c, ["99490", "99491"], None, 1)
        assert diffs == ()

    def test_only_the_newest_year_per_code_is_compared(self, con):
        c, _ = con
        write_elements(
            c,
            "99490",
            [
                _el("99490", 2020, "consent", "verbal"),
                _el("99490", 2026, "consent", "required"),
            ],
            [],
        )
        write_elements(c, "99491", [_el("99491", 2026, "consent", "verbal")], [])
        diffs, _ = lineage._element_diffs(c, ["99490", "99491"], None, 1)
        # 2020's "verbal" for 99490 must not be compared — only 2026's
        # "required" — so the diff is required (99490) vs verbal (99491).
        values = {d.value: d for d in diffs}
        assert "verbal" not in {d.value for d in diffs if "99490" in d.in_codes}
        assert values["required"].in_codes == ("99490",)

    def test_capped_at_twelve_most_shared_first(self, con):
        c, _ = con
        codes = ["A0001", "A0002", "A0003"]
        for code in codes:
            rows = [_el(code, 2026, f"t{i}", f"v{i}") for i in range(20)]
            # one shared value across all three codes, sorts first
            rows.append(_el(code, 2026, "shared", "yes"))
            write_elements(c, code, rows, [])
        # a fourth code with none of the "t*" values, so every t* pair
        # differs (present in the three, absent in the fourth)
        write_elements(c, "A0004", [_el("A0004", 2026, "shared", "no")], [])
        diffs, _ = lineage._element_diffs(c, codes + ["A0004"], None, 1)
        assert len(diffs) == 12
        assert diffs[0].in_codes == ("A0001", "A0002", "A0003")

    def test_element_label_cpt_source_with_no_anchor_falls_back_to_cpt_changes(self):
        """``_element_label`` prefers a resolved FR paragraph label; with
        neither ``item_key`` nor ``p_id`` and a CPT-codebook source, it
        falls back to "CPT Changes {year}" (mirrors ``_to_lineage_event``'s
        own cpt fallback, tested in ``TestCptLabel``)."""
        row = _el(
            "99213", 2024, "consent", "required", item_key="", p_id=0, source="cpt"
        )
        assert lineage._element_label(None, row, 1) == "CPT Changes 2024"


class TestElementDiffsErrors:
    def test_missing_table_is_swallowed_as_no_elements(self):
        bare = duckdb.connect(":memory:")  # pfs.code_element doesn't exist
        assert lineage._element_diffs(bare, ["99490", "99491"], None, 1) == ((), "")

    def test_a_real_error_propagates(self):
        with patch("llm.lineage.read_elements", side_effect=ValueError("boom")):
            with pytest.raises(ValueError):
                lineage._element_diffs(None, ["99490"], None, 1)


class TestGuidance:
    def test_cfr_first_deduped_and_capped(self, con, store):
        c, _ = con
        rows = [
            _gd("CCM", "99490", "cfr", f"42 CFR 410.{i}", "YBM4IZUS", 686)
            for i in range(6)
        ]
        rows += [
            _gd("CCM", "99490", "iom", f"100-04 ch.{i}", "YBM4IZUS", 686)
            for i in range(6)
        ]
        write_guidance(c, "CCM", rows)
        out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
        assert len(out) == 8
        assert [g.kind for g in out[:6]] == ["cfr"] * 6
        assert out[0].url == "https://www.ecfr.gov/current/title-42/section-410.0"

    def test_dedupes_by_locator(self, con, store):
        c, _ = con
        write_guidance(
            c,
            "CCM",
            [
                _gd("CCM", "99490", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 1578),
                _gd("CCM", "99491", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 686),
            ],
        )
        out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
        assert len(out) == 1

    def test_wide_families_are_skipped(self, con, store):
        c, _ = con
        write_guidance(
            c, "APCM", [_gd("APCM", "G0556", "mln", "MLN 907166", "YBM4IZUS", 686)]
        )
        out = lineage._collect_guidance(c, store, ["APCM"], ["APCM"], 1)
        assert out == ()

    def test_iom_mln_have_no_url(self, con, store):
        c, _ = con
        write_guidance(
            c,
            "CCM",
            [_gd("CCM", "99490", "iom", "100-04 ch.12 §30.6.4", "YBM4IZUS", 686)],
        )
        out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
        assert out[0].url == ""
        assert out[0].label == "CY2021 PFS final 85 FR 84547 ¶686"

    def test_missing_table_is_swallowed_as_no_guidance(self):
        bare = duckdb.connect(":memory:")  # pfs.code_guidance doesn't exist
        assert lineage._collect_guidance(bare, None, ["CCM"], [], 1) == ()

    def test_a_real_error_propagates(self):
        with patch("llm.lineage.read_guidance", side_effect=ValueError("boom")):
            with pytest.raises(ValueError):
                lineage._collect_guidance(None, None, ["CCM"], [], 1)

    def test_cfr_url_unresolved_degrades_to_empty_string_not_dropped(self, con, store):
        """A CFR locator that ``bib.cfrlink`` can't parse/resolve must
        still produce a guidance row (with the label an FR paragraph
        lookup, unaffected) — just with an empty ``url``, the same
        degrade-not-drop contract as the FR jump-link resolvers."""
        c, _ = con
        write_guidance(
            c,
            "CCM",
            [_gd("CCM", "99490", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 686)],
        )
        with patch("bib.cfrlink.parse_cite", side_effect=RuntimeError("boom")):
            out = lineage._collect_guidance(c, store, ["CCM"], [], 1)
        assert len(out) == 1
        assert out[0].url == ""
        assert out[0].locator == "42 CFR 410.78(a)(3)"


class TestCodesStr:
    """Ruling B11: element-diff code lists compact past 8 codes."""

    def test_short_list_joined_plainly(self):
        assert lineage._codes_str(("A", "B", "C")) == "A, B, C"

    def test_exactly_eight_stays_plain(self):
        codes = tuple(f"C{i}" for i in range(8))
        assert lineage._codes_str(codes) == ", ".join(codes)

    def test_over_eight_compacts_to_count_and_first_three(self):
        codes = tuple(f"C{i}" for i in range(9))
        assert lineage._codes_str(codes) == "9 codes (C0, C1, C2, …)"


class TestPromptBlock:
    def test_exact_text(self):
        events = (
            LineageEvent(
                code="G2058",
                year=2021,
                kind="replaced_by",
                from_codes=(),
                to_codes=("99439",),
                label="CY2021 PFS final ¶686",
                item_key="YBM4IZUS",
                p_id=686,
                page=84547,
                url="https://fr.test/2021-doc#p-686",
                source="fr",
                anchored=True,
                note="",
            ),
            LineageEvent(
                code="99439",
                year=2021,
                kind="created",
                from_codes=(),
                to_codes=(),
                label="CY2021 PFS final ¶686",
                item_key="YBM4IZUS",
                p_id=686,
                page=84547,
                url="https://fr.test/2021-doc#p-686",
                source="fr",
                anchored=True,
                note="from CPT Editorial Panel",
            ),
        )
        diffs = (
            ElementDiff(
                type="consent",
                value="required",
                in_codes=("99490", "99491"),
                not_in_codes=("G0556",),
                label="CY2021 PFS final ¶1578",
                item_key="YBM4IZUS",
                p_id=1578,
            ),
        )
        guidance = (
            GuidanceRef(
                kind="cfr",
                locator="42 CFR 410.78(a)(3)",
                url="https://www.ecfr.gov/current/title-42/section-410.78",
                label="CY2021 PFS final ¶1578",
                item_key_src="YBM4IZUS",
                p_id_src=1578,
            ),
        )
        ev = LineageEvidence(
            codes=("99439", "99490", "99491", "G0556", "G2058"),
            families=("CCM", "APCM"),
            events=events,
            element_diffs=diffs,
            guidance=guidance,
        )
        assert ev.prompt_block() == (
            lineage._HEADER + "\n"
            "2021: G2058 was replaced by 99439 [CY2021 PFS final ¶686]\n"
            "2021: 99439 was created — from CPT Editorial Panel "
            "[CY2021 PFS final ¶686]\n"
            "Element differences:\n"
            "consent=required: required for 99490, 99491; not for G0556 "
            "[CY2021 PFS final ¶1578]\n"
            "Guidance:\n"
            "42 CFR 410.78(a)(3) (CFR) is cited for these codes "
            "[CY2021 PFS final ¶1578]"
        )

    def test_header_only_with_no_events(self):
        ev = LineageEvidence(
            codes=("99490",), families=(), events=(), element_diffs=(), guidance=()
        )
        assert ev.prompt_block() == lineage._HEADER

    def test_event_lines_are_sentences_ending_with_the_label(self):
        """qwen2.5 read the old "[label] 2021 replaced_by G2058 ( → 99439)"
        rows as a table it did not have to cite; a sentence ending in the
        label is the shape the system prompt asks the answer to copy."""

        def line(kind, frm=(), to=(), note=""):
            return lineage._event_line(
                LineageEvent(
                    code="G2058",
                    year=2021,
                    kind=kind,
                    from_codes=frm,
                    to_codes=to,
                    label="L",
                    item_key="K",
                    p_id=1,
                    page=1,
                    url="",
                    source="fr",
                    anchored=True,
                    note=note,
                )
            )

        assert (
            line("replaced_by", to=("99439",))
            == "2021: G2058 was replaced by 99439 [L]"
        )
        assert line("replaces", frm=("G0506",)) == "2021: G2058 replaced G0506 [L]"
        assert line("created") == "2021: G2058 was created [L]"
        assert (
            line("crosswalk", to=("99439",))
            == "2021: G2058 was crosswalked to 99439 [L]"
        )
        assert line("disappeared") == "2021: G2058 dropped out of the PFS RVU file [L]"
        # a directional kind whose codes were filtered away still reads
        assert line("replaced_by") == "2021: G2058 was replaced by [L]"

    def test_from_and_to_appended_as_suffixes_when_the_template_has_no_placeholder(
        self,
    ):
        """``status_change``'s template ("changed payment status") embeds
        neither ``{from_}`` nor ``{to}`` — non-empty from/to codes must
        still reach the model, as trailing "(from …)"/"(to …)" clauses."""
        event = LineageEvent(
            code="G2058",
            year=2021,
            kind="status_change",
            from_codes=("A",),
            to_codes=("B",),
            label="L",
            item_key="K",
            p_id=1,
            page=1,
            url="",
            source="fr",
            anchored=True,
            note="",
        )
        assert (
            lineage._event_line(event)
            == "2021: G2058 changed payment status (from A) (to B) [L]"
        )

    def test_priority_kinds_exceed_cap_but_are_never_dropped(self):
        priorities = [
            LineageEvent(
                "99490", 2020 + i, k, (), (), f"[L{i}]", "K", 1, 0, "", "fr", True, ""
            )
            for i, k in enumerate(["created", "replaces", "replaced_by", "deleted"])
        ]
        ev = LineageEvidence(
            codes=("99490",),
            families=(),
            events=tuple(priorities),
            element_diffs=(),
            guidance=(),
            max_prompt_rows=2,
        )
        rendered = ev.prompt_block()
        for i in range(4):  # every priority row rendered, each ending with its label
            assert f"[[L{i}]]" in rendered

    def test_under_budget_every_priority_row_for_a_code_survives(self):
        """Ruling B15: the per-code priority cap is a budget safety net,
        not a default trim — under the hard cap (2 * max_prompt_rows),
        a single code's 6 legitimate priority-kind rows (e.g. several
        genuine replaces/replaced_by events over the years) all survive,
        exactly as before B11 introduced the per-code cap."""
        priorities = [
            LineageEvent(
                "99490",
                2020 + i,
                "replaces",
                (),
                (),
                f"[L{i}]",
                "K",
                i,
                0,
                "",
                "fr",
                True,
                "",
            )
            for i in range(6)
        ]
        ev = LineageEvidence(
            codes=("99490",),
            families=(),
            events=tuple(priorities),
            element_diffs=(),
            guidance=(),
            max_prompt_rows=25,  # hard cap 50 — 6 priority rows fit easily
        )
        rendered = ev.prompt_block()
        assert sum(rendered.count(f"[L{i}]") for i in range(6)) == 6

    def test_over_budget_priority_rows_for_a_code_trim_to_four(self):
        """Ruling B15: the same 6 priority-kind rows for one code, but
        with a hard cap (2 * max_prompt_rows) too small to hold them —
        the per-code cap of 4 engages as the safety net."""
        priorities = [
            LineageEvent(
                "99490",
                2020 + i,
                "replaces",
                (),
                (),
                f"[L{i}]",
                "K",
                i,
                0,
                "",
                "fr",
                True,
                "",
            )
            for i in range(6)
        ]
        ev = LineageEvidence(
            codes=("99490",),
            families=(),
            events=tuple(priorities),
            element_diffs=(),
            guidance=(),
            max_prompt_rows=2,  # hard cap 4 — 6 priority rows don't fit
        )
        rendered = ev.prompt_block()
        assert sum(rendered.count(f"[L{i}]") for i in range(6)) == 4

    def test_hard_cap_at_twice_max_rows(self):
        """Ruling B11: many codes each with priority rows can push the
        priority selection itself past max_rows — the absolute ceiling
        is 2 * max_rows."""
        priorities = [
            LineageEvent(
                f"C{i}",
                2020,
                "created",
                (),
                (),
                f"[L{i}]",
                "K",
                i,
                0,
                "",
                "fr",
                True,
                "",
            )
            for i in range(20)  # 20 distinct codes, one priority row each
        ]
        ev = LineageEvidence(
            codes=tuple(f"C{i}" for i in range(20)),
            families=(),
            events=tuple(priorities),
            element_diffs=(),
            guidance=(),
            max_prompt_rows=5,
        )
        rendered = ev.prompt_block()
        kept = sum(rendered.count(f"[L{i}]") for i in range(20))
        assert kept == 10  # 2 * max_prompt_rows(5)


class TestPayload:
    def test_json_serializable_and_shaped(self):
        events = (
            LineageEvent(
                "99439",
                2021,
                "created",
                (),
                (),
                "[L]",
                "YBM4IZUS",
                686,
                84547,
                "u",
                "fr",
                True,
                "",
            ),
        )
        ev = LineageEvidence(
            codes=("99439",),
            families=("CCM",),
            events=events,
            element_diffs=(),
            guidance=(),
            elements_note="elements extracted for 1 of 3 codes",
        )
        payload = ev.payload()
        raw = json.dumps(payload)
        back = json.loads(raw)
        assert back["type"] == "lineage"
        assert back["codes"] == ["99439"]
        assert back["families"] == ["CCM"]
        assert back["events"][0]["kind"] == "created"
        assert back["elements_note"] == "elements extracted for 1 of 3 codes"

    def test_no_elements_note_key_when_empty(self):
        ev = LineageEvidence(("99439",), (), (), (), ())
        assert "elements_note" not in ev.payload()

    def test_explicit_key_present_when_non_empty(self):
        """Ruling B10: codes literally named in the question ride along
        in the payload so a client could highlight them — omitted
        entirely (not an empty list) when nothing was named explicitly,
        covered by ``test_no_elements_note_key_when_empty``'s sibling
        default-``LineageEvidence`` above."""
        ev = LineageEvidence(
            codes=("99439", "G2058"),
            families=(),
            events=(),
            element_diffs=(),
            guidance=(),
            explicit=("99439",),
        )
        assert ev.payload()["explicit"] == ["99439"]

    def test_no_explicit_key_when_empty(self):
        ev = LineageEvidence(("99439",), (), (), (), (), explicit=())
        assert "explicit" not in ev.payload()


class TestLineageEvidence:
    def test_none_when_no_codes_detected(self):
        assert (
            lineage_evidence("what did commenters say about telehealth?", CFG) is None
        )

    def test_collapsed_events_and_labels_for_the_ccm_fixture(self, con, store):
        c, path = con
        write_events(
            c,
            "99490",
            [
                _ev(
                    "99490",
                    2015,
                    "adopted_cpt",
                    item_key="DE2VH9PD",
                    p_id=1251,
                    page=67716,
                )
            ],
        )
        write_events(
            c,
            "G2058",
            [
                _ev(
                    "G2058",
                    2021,
                    "replaced_by",
                    to="99439",
                    item_key="YBM4IZUS",
                    p_id=686,
                    page=84547,
                ),
                _ev("G2058", 2022, "disappeared", source="rvu", item_key="", p_id=0),
            ],
        )
        write_events(
            c,
            "99439",
            [
                _ev(
                    "99439", 2021, "created", item_key="YBM4IZUS", p_id=686, page=84547
                ),
                _ev(
                    "99439",
                    2021,
                    "replaces",
                    frm="G2058",
                    item_key="YBM4IZUS",
                    p_id=1578,
                    page=84639,
                ),
                # same group as the fr "created" row above — collapse must
                # drop this rvu duplicate.
                _ev(
                    "99439",
                    2021,
                    "created",
                    source="rvu",
                    item_key="",
                    p_id=0,
                    note="dupe",
                ),
            ],
        )
        c.close()  # release the write handle before lineage_evidence opens read-only
        cfg = replace(CFG, duckdb_replica=str(path))
        det = Detection(
            codes=("99439", "99490", "G2058"), families=(), explicit=(), wide=()
        )
        with patch("llm.lineage._store", return_value=store):
            with patch("llm.lineage.detect_codes", return_value=det):
                ev = lineage_evidence("irrelevant text", cfg)
        assert ev is not None
        assert [(e.year, e.code, e.kind, e.source) for e in ev.events] == [
            (2015, "99490", "adopted_cpt", "fr"),
            (2021, "99439", "created", "fr"),
            (2021, "99439", "replaces", "fr"),
            (2021, "G2058", "replaced_by", "fr"),
            (2022, "G2058", "disappeared", "rvu"),
        ]
        adopted, created, replaces, replaced_by, disappeared = ev.events
        assert adopted.label == "CY2015 PFS final 79 FR 67716 ¶1251"
        assert created.label == "CY2021 PFS final 85 FR 84547 ¶686"
        assert replaces.label == "CY2021 PFS final 85 FR 84639 ¶1578"
        assert replaces.from_codes == ("G2058",)
        assert replaced_by.to_codes == ("99439",)
        assert disappeared.label == "PFS CY2022 RVU file"
        assert created.url == "https://fr.test/2021-doc#p-686"

    def test_on_demand_fallback_only_for_codes_without_rows_and_capped(
        self, con, store
    ):
        c, path = con
        rvu_cols = (
            "hcpcs VARCHAR, mod VARCHAR, description VARCHAR, status_code VARCHAR, "
            "non_fac_total DOUBLE, year INTEGER"
        )
        c.execute(f"CREATE TABLE pfs.rvu ({rvu_cols})")
        c.executemany(
            "INSERT INTO pfs.rvu VALUES (?,?,?,?,?,?)",
            [
                *[("00000", None, "filler", "A", 1.0, y) for y in range(2015, 2027)],
                ("99441", None, "d", "A", 1.0, 2020),
                ("99441", None, "d", "A", 1.0, 2021),
                ("99442", None, "d", "A", 1.0, 2020),
                ("99442", None, "d", "A", 1.0, 2021),
            ],
        )
        # 99490 has precomputed rows — must not trigger the fallback.
        write_events(
            c,
            "99490",
            [_ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251)],
        )
        c.close()  # release the write handle before lineage_evidence opens read-only
        cfg = replace(CFG, duckdb_replica=str(path), lineage_on_demand_max=1)
        det = Detection(
            codes=("99441", "99442", "99490"), families=(), explicit=(), wide=()
        )
        with patch("llm.lineage._store", return_value=store):
            with patch("llm.lineage.detect_codes", return_value=det):
                with patch(
                    "llm.lineage._on_demand_events", wraps=lineage._on_demand_events
                ) as spy:
                    ev = lineage_evidence("irrelevant text", cfg)
        assert ev is not None
        # budget is 1: only the first no-rows code (99441, sorted first)
        # gets the fallback; 99442 and the precomputed 99490 do not.
        assert spy.call_count == 1
        assert spy.call_args.args[-1] == "99441"
        codes_with_events = {e.code for e in ev.events}
        assert "99441" in codes_with_events
        assert "99442" not in codes_with_events

    def test_never_raises_on_a_broken_replica(self, tmp_path):
        bad = tmp_path / "not-a-duckdb-file"
        bad.write_text("not a database")
        cfg = replace(CFG, duckdb_replica=str(bad))
        det = Detection(codes=("99490",), families=(), explicit=(), wide=())
        with patch("llm.lineage.detect_codes", return_value=det):
            assert lineage_evidence("irrelevant", cfg) is None

    def test_elements_note_and_guidance_flow_through(self, con, store):
        c, path = con
        write_elements(c, "99490", [_el("99490", 2026, "consent", "required")], [])
        write_guidance(
            c,
            "CCM",
            [_gd("CCM", "99490", "cfr", "42 CFR 410.78(a)(3)", "YBM4IZUS", 1578)],
        )
        c.close()  # release the write handle before lineage_evidence opens read-only
        cfg = replace(CFG, duckdb_replica=str(path))
        det = Detection(codes=("99490",), families=("CCM",), explicit=(), wide=())
        with patch("llm.lineage._store", return_value=store):
            with patch("llm.lineage.detect_codes", return_value=det):
                ev = lineage_evidence("irrelevant", cfg)
        assert ev is not None
        assert ev.elements_note == "elements extracted for 1 of 1 codes"
        assert ev.element_diffs == ()
        assert len(ev.guidance) == 1
        assert ev.guidance[0].locator == "42 CFR 410.78(a)(3)"

    def test_never_raises_when_a_query_step_fails_unexpectedly(self, con, store):
        """The replica opens fine and the codes are detected, but
        something inside the query try-block breaks unexpectedly (not a
        missing-table/broken-replica case, which have their own
        coverage) — ``lineage_evidence`` must still degrade to ``None``,
        not propagate."""
        c, path = con
        c.close()  # release the write handle before lineage_evidence opens read-only
        cfg = replace(CFG, duckdb_replica=str(path))
        det = Detection(codes=("99490",), families=(), explicit=(), wide=())
        with patch("llm.lineage._store", return_value=store):
            with patch("llm.lineage.detect_codes", return_value=det):
                with patch(
                    "llm.lineage._collect_events", side_effect=RuntimeError("boom")
                ):
                    assert lineage_evidence("irrelevant", cfg) is None


class TestCptLabel:
    def test_cpt_label_is_always_cpt_changes_year_note_holds_the_detail(self):
        """CPT-sourced events (cpt_changed, cpt_deleted) always label as
        "CPT Changes {year}" — the note (e.g. "absent from CPT 2024")
        is a separate field, not folded into the label."""
        row = _ev(
            "99213",
            2024,
            "cpt_deleted",
            item_key="CPTED2024",
            p_id=0,
            source="cpt",
            note="absent from CPT 2024",
        )
        event = lineage._to_lineage_event(None, row, 1)
        assert event.label == "CPT Changes 2024"
        assert event.note == "absent from CPT 2024"
        line = lineage._event_line(event)
        assert (
            line
            == "2024: 99213 was deleted from CPT — absent from CPT 2024 [CPT Changes 2024]"
        )
        assert (
            line.count("absent from CPT 2024") == 1
        )  # not duplicated into the label too

    def test_cpt_changed_label_ignores_any_note(self):
        row = _ev(
            "99213",
            2022,
            "cpt_changed",
            item_key="CPTED2022",
            p_id=0,
            source="cpt",
            note="some other detail",
        )
        event = lineage._to_lineage_event(None, row, 1)
        assert event.label == "CPT Changes 2022"


class TestSourceLabelAndUrlDegrade:
    def test_source_label_degrades_to_item_key_when_store_get_fails(self):
        store = SimpleNamespace(get=lambda k: (_ for _ in ()).throw(KeyError(k)))
        assert lineage._source_label(store, "UNKNOWN1", 5, 1) == "UNKNOWN1"

    def test_source_label_degrades_to_question_mark_when_item_key_empty(self):
        store = SimpleNamespace(get=lambda k: (_ for _ in ()).throw(KeyError(k)))
        assert lineage._source_label(store, "", 5, 1) == "?"

    def test_fr_url_degrades_to_empty_string_on_resolve_failure(self, store):
        with patch("bib.frlink.resolve", side_effect=RuntimeError("boom")):
            assert lineage._fr_url(store, "YBM4IZUS", 686, 1) == ""

    def test_fr_url_empty_when_item_key_or_p_id_missing(self):
        assert lineage._fr_url(None, "", 5, 1) == ""
        assert lineage._fr_url(None, "K", 0, 1) == ""

    def test_rule_kind_degrades_to_rule_when_bib_frlink_raises(self, store):
        """``_rule_kind`` must never say ``final``/``proposed`` on a
        failure path — a label falsely claiming a rule's kind is a false
        statement of policy (ruling in ``rule_label``'s docstring)."""
        with patch("bib.frlink.rule_kind", side_effect=RuntimeError("boom")):
            kind = lineage._rule_kind(store, "YBM4IZUS", 918_273)
        assert kind == "rule"


class TestCachingAndScopedUrlResolution:
    """Ruling: Store.get results are cached per item_key (not just per
    (item_key, p_id) — the same rule backs many events/diffs/guidance
    rows), and FR urls are only resolved for the rows that actually
    reach the prompt block."""

    def test_urls_resolved_only_for_prompt_rows_labels_cached_per_item(
        self, con, store
    ):
        c, path = con
        n_items = 10
        rows = [
            _ev(
                "99999",
                2000 + i,
                "crosswalk",  # non-priority, not an _EARLIEST_ONLY kind
                item_key=f"K{i % n_items:02d}",
                p_id=1000 + i,
                source="fr",
                anchored=True,
            )
            for i in range(60)
        ]
        write_events(c, "99999", rows)
        store.con.executemany(
            "INSERT INTO items VALUES (?,?,?,?)",
            [
                (
                    f"K{i:02d}",
                    f"Medicare Program; CY {2000 + i} PFS Final Rule",
                    f"{1999 + i}-11-01",
                    f"https://www.federalregister.gov/d/{2000 + i}-doc",
                )
                for i in range(n_items)
            ],
        )
        c.close()  # release the write handle before lineage_evidence opens read-only
        cfg = replace(CFG, duckdb_replica=str(path), lineage_max_rows=25)
        det = Detection(codes=("99999",), families=(), explicit=(), wide=())
        real_get = store.get
        seen_keys: list[str] = []

        def counting_get(key):
            seen_keys.append(key)
            return real_get(key)

        with patch("llm.lineage._store", return_value=store):
            with patch("llm.lineage.detect_codes", return_value=det):
                with patch.object(store, "get", side_effect=counting_get) as get_spy:
                    with patch("bib.frlink.resolve") as resolve_mock:
                        resolve_mock.return_value = SimpleNamespace(url="http://x")
                        ev = lineage_evidence("irrelevant", cfg)

        assert ev is not None
        assert len(ev.events) == 60  # payload carries every collapsed row
        # exactly the 25 prompt-selected rows (budget: no priority kinds
        # here, so the first 25 by year) get an FR url resolved.
        assert resolve_mock.call_count == 25
        with_url = [e for e in ev.events if e.url]
        assert len(with_url) == 25
        # every referenced item_key is fetched at most once, total calls
        # bounded by the number of distinct item_keys (10), not events (60).
        assert get_spy.call_count <= n_items
        assert len(seen_keys) == len(set(seen_keys))  # never re-fetched


class TestOnDemandEventsCache:
    def test_second_call_with_the_same_code_and_mtime_hits_the_cache(self):
        calls: list[str] = []

        def fake_lineage(cur, store, code):
            calls.append(code)
            return [_ev(code, 2021, "created")]

        with patch("pfs.lineage.lineage", side_effect=fake_lineage):
            first = lineage._on_demand_events(None, None, 55, "99490")
            second = lineage._on_demand_events(None, None, 55, "99490")
        assert calls == ["99490"]  # the underlying lookup ran only once
        assert first == second == [_ev("99490", 2021, "created")]


class TestCollectEvents:
    def test_missing_table_is_swallowed_as_no_rows(self):
        bare = duckdb.connect(":memory:")  # pfs.code_event doesn't exist
        assert lineage._collect_events(bare, None, 0, ["99490"], 0) == []

    def test_a_real_error_propagates(self):
        with patch("llm.lineage.read_events", side_effect=ValueError("boom")):
            with pytest.raises(ValueError):
                lineage._collect_events(None, None, 0, ["99490"], 0)


class TestConcurrentLineage:
    """C1/Ruling B14 — ``llm.lineage._store()`` must be thread-local:
    bib.Store's sqlite connection is check_same_thread=True, and /chat
    runs each turn in Starlette's threadpool. Mirrors
    ``TestConcurrentChats`` in test_evidence.py."""

    @pytest.fixture
    def bib_db(self, tmp_path):
        from bib.item import Rule
        from bib.store import Store

        db = tmp_path / "bib.sqlite"
        seed = Store(str(db), storage_dir=str(tmp_path / "storage"))
        key = seed.create(
            Rule(
                title="Medicare Program; CY 2021 Payment Policies Under the PFS",
                date_published="2020-12-28",
            )
        )
        con = seed._con()
        con.execute(
            "INSERT INTO fr_anchor_docs "
            "(item_key, document_number, html_url, start_page, end_page, fr_volume) "
            "VALUES (?,?,?,?,?,?)",
            (key, "2020-2021doc", "https://fr.test/2021-doc", 84000, 85000, 85),
        )
        con.execute(
            "INSERT INTO fr_anchors (item_key, p_id, page, ordinal, text) "
            "VALUES (?,?,?,?,?)",
            (
                key,
                686,
                84547,
                1,
                "We are finalizing HCPCS code G2058 as new CPT code 99439.",
            ),
        )
        con.commit()
        seed.close()
        return str(db), key

    def test_no_cross_thread_sqlite_errors_one_store_per_thread(
        self, bib_db, monkeypatch
    ):
        from bib.store import Store

        db_path, item_key = bib_db
        opens: list[int] = []
        real_init = Store.__init__

        def counting_init(self, *a, **kw):
            opens.append(threading.get_ident())
            real_init(self, db_path)

        monkeypatch.setattr(Store, "__init__", counting_init)
        monkeypatch.setattr("conf.connect.bib", lambda: Store())
        monkeypatch.setattr(
            "bib.frlink.resolve",
            lambda ref, *, store, item_key: SimpleNamespace(
                url=f"https://fr.test/{item_key}#{ref}"
            ),
        )
        # I2 (later in this fix wave) adds "{vol} FR {page}" to this label
        # once volume/page are resolvable — DE2VH9PD/YBM4IZUS-style fixture
        # data does resolve them, so this constant moves with that change.
        expected_label = "CY2021 PFS final 85 FR 84547 ¶686"

        def event() -> LineageEvent:
            store = lineage._store()
            return lineage._to_lineage_event(
                store, _ev("99439", 2021, "created", item_key=item_key, p_id=686), 1
            )

        errors: list[BaseException] = []
        n_calls = 40

        def worker(_i: int) -> None:
            try:
                for _ in range(n_calls):
                    store = lineage._store()
                    label = lineage._source_label(store, item_key, 686, 1)
                    url = lineage._fr_url(store, item_key, 686, 1)
                    assert label == expected_label, label
                    assert url == f"https://fr.test/{item_key}#p-686", url
                    ev = _lineage_evidence([event()])
                    (src,) = lineage.lineage_sources(store, ev, max_items=6)
                    assert src["label"] == expected_label
            except BaseException as e:  # noqa: BLE001 — collected, not raised in-thread
                errors.append(e)

        with ThreadPoolExecutor(max_workers=8) as pool:
            list(pool.map(worker, range(8)))

        assert errors == [], errors
        # one Store construction per thread (threading.local caching),
        # not one per call (320) and not one shared across every thread.
        assert len(opens) == 8
        assert len(set(opens)) == 8


def _lineage_evidence(
    events, max_prompt_rows=25, explicit=(), element_diffs=()
) -> LineageEvidence:
    return LineageEvidence(
        codes=(),
        families=(),
        events=tuple(events),
        element_diffs=tuple(element_diffs),
        guidance=(),
        max_prompt_rows=max_prompt_rows,
        explicit=tuple(explicit),
    )


class TestLineageSources:
    """``lineage_sources`` (Ruling B6) — FR paragraphs behind the
    events selected into the prompt block, as rule-kind sources whose
    label is the lineage event's own label."""

    def test_priority_kinds_first_then_by_year(self, store):
        # revalued (non-priority) at 2020 must still sort after every
        # priority-kind event, even one dated later (2021).
        revalued = lineage._to_lineage_event(
            store,
            _ev("99439", 2020, "revalued", item_key="YBM4IZUS", p_id=1578),
            mtime=0,
        )
        created = lineage._to_lineage_event(
            store,
            _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        adopted = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        ev = _lineage_evidence([revalued, created, adopted])
        out = lineage_sources(store, ev, max_items=6)
        assert [s["label"] for s in out] == [
            adopted.label,
            created.label,
            revalued.label,
        ]

    def test_capped_at_max_items(self, store):
        store.con.execute(
            "INSERT INTO items VALUES (?,?,?,?)",
            (
                "EXTRA1",
                "Medicare Program; CY 2022 PFS Final Rule",
                "2021-11-02",
                "https://www.federalregister.gov/d/2021-2022doc",
            ),
        )
        store.con.execute(
            "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
            ("EXTRA1", 10, 1000, 1, "extra paragraph text"),
        )
        store.con.commit()
        events = [
            lineage._to_lineage_event(
                store,
                _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
                mtime=0,
            ),
            lineage._to_lineage_event(
                store,
                _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686),
                mtime=0,
            ),
            lineage._to_lineage_event(
                store,
                _ev(
                    "99439",
                    2021,
                    "replaces",
                    frm="G2058",
                    item_key="YBM4IZUS",
                    p_id=1578,
                ),
                mtime=0,
            ),
            lineage._to_lineage_event(
                store,
                _ev("G0000", 2022, "revalued", item_key="EXTRA1", p_id=10),
                mtime=0,
            ),
        ]
        ev = _lineage_evidence(events)
        out = lineage_sources(store, ev, max_items=2)
        assert len(out) == 2

    def test_event_loop_enforces_the_hard_cap_even_with_a_raised_max_items(self, store):
        """Ruling B11: _LINEAGE_SOURCES_HARD_CAP (10) is an absolute
        ceiling the event loop must enforce itself — a config-raised
        max_items must not let it through."""
        events = []
        for i in range(15):
            key = f"MANY{i}"
            _insert_extra_anchor(store, key, 2000 + i, 100 + i)
            events.append(
                lineage._to_lineage_event(
                    store,
                    _ev(f"M{i}", 2000 + i, "revalued", item_key=key, p_id=100 + i),
                    mtime=0,
                )
            )
        ev = _lineage_evidence(events, max_prompt_rows=25)
        out = lineage_sources(store, ev, max_items=20)
        assert len(out) == 10

    def test_label_equals_the_event_label(self, store):
        e = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        ev = _lineage_evidence([e])
        (out,) = lineage_sources(store, ev, max_items=6)
        assert out["label"] == e.label == "CY2015 PFS final 79 FR 67716 ¶1251"
        assert out["id"] == e.label
        assert out["kind"] == "rule"
        assert out["item_key"] == "DE2VH9PD"
        assert out["p_id"] == "1251"
        assert "adopt CPT code 99490" in out["snippet"]

    def test_dedupes_distinct_item_key_p_id_across_events(self, store):
        # two collapsed events can point at the same paragraph (e.g. a
        # "created" and a "replaces" row both anchored at YBM4IZUS 686
        # in a contrived case) — only one source per (item_key, p_id).
        e1 = lineage._to_lineage_event(
            store,
            _ev("99439", 2021, "created", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        e2 = lineage._to_lineage_event(
            store,
            _ev("99439", 2021, "revalued", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        ev = _lineage_evidence([e1, e2])
        out = lineage_sources(store, ev, max_items=6)
        assert len(out) == 1

    def test_skips_non_fr_events(self, store):
        rvu = lineage._to_lineage_event(
            store, _ev("99441", 2022, "disappeared", source="rvu"), mtime=0
        )
        ev = _lineage_evidence([rvu])
        assert lineage_sources(store, ev, max_items=6) == []

    def test_missing_paragraph_is_skipped_not_raised(self, store):
        e = lineage._to_lineage_event(
            store,
            _ev("99999", 2019, "created", item_key="NOPE", p_id=9999),
            mtime=0,
        )
        ev = _lineage_evidence([e])
        assert lineage_sources(store, ev, max_items=6) == []

    def test_never_raises_on_a_broken_store(self):
        class _BrokenStore:
            def _con(self):
                raise RuntimeError("boom")

        e = LineageEvent(
            code="99490",
            year=2015,
            kind="adopted_cpt",
            from_codes=(),
            to_codes=(),
            label="CY2015 PFS final ¶1251",
            item_key="DE2VH9PD",
            p_id=1251,
            page=67716,
            url="",
            source="fr",
            anchored=True,
            note="",
        )
        ev = _lineage_evidence([e])
        assert lineage_sources(_BrokenStore(), ev, max_items=6) == []

    def test_only_prompt_selected_events_are_considered(self, store):
        # max_prompt_rows=0 with a non-priority event only — priority
        # rows are never dropped by _select_for_prompt, but a
        # non-priority row past budget zero is.
        e = lineage._to_lineage_event(
            store,
            _ev("99439", 2021, "revalued", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        ev = _lineage_evidence([e], max_prompt_rows=0)
        assert lineage_sources(store, ev, max_items=6) == []

    def test_never_raises_when_selecting_prompt_events_fails(self, store):
        """A failure before the candidate loop even starts (selecting/
        sorting the prompt events) must degrade to ``[]``, the same
        contract ``lineage_evidence`` itself has — never break the chat
        over a sources lookup."""
        e = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        ev = _lineage_evidence([e])
        with patch("llm.lineage._select_for_prompt", side_effect=RuntimeError("boom")):
            assert lineage_sources(store, ev, max_items=6) == []


def _insert_extra_anchor(store, key, year, p_id, text="extra text"):
    store.con.execute(
        "INSERT INTO items VALUES (?,?,?,?)",
        (
            key,
            f"Medicare Program; CY {year} PFS Final Rule",
            f"{year - 1}-11-01",
            f"https://www.federalregister.gov/d/{year}-doc-{key}",
        ),
    )
    store.con.execute(
        "INSERT INTO fr_anchors VALUES (?,?,?,?,?)",
        (key, p_id, 1000 + p_id, 1, text),
    )
    store.con.commit()


def _diff(item_key, p_id, label, type_="element", value="v") -> ElementDiff:
    return ElementDiff(
        type=type_,
        value=value,
        in_codes=("A",),
        not_in_codes=("B",),
        label=label,
        item_key=item_key,
        p_id=p_id,
    )


class TestLineageSourcesRulingB10:
    """Ruling B10 — explicit-code-first ordering and element-diff
    anchors appended to ``lineage_sources``' output."""

    def test_explicit_code_sorts_first_even_when_oldest_and_non_priority(self, store):
        # A (explicit, literally named in the question) has the oldest,
        # non-priority-kind event; B and C (family-mates, not named) have
        # newer, priority-kind events. Explicit-code membership must
        # outrank both the priority-kind tier and the year — A still
        # sorts first.
        a = lineage._to_lineage_event(
            store,
            _ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        b = lineage._to_lineage_event(
            store,
            _ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        c = lineage._to_lineage_event(
            store,
            _ev("C", 2021, "created", item_key="YBM4IZUS", p_id=1578),
            mtime=0,
        )
        ev = _lineage_evidence([a, b, c], explicit=("A",))
        out = lineage_sources(store, ev, max_items=8)
        assert [s["label"] for s in out] == [a.label, b.label, c.label]

    def test_explicit_code_survives_a_cap_that_would_otherwise_drop_it(self, store):
        a = lineage._to_lineage_event(
            store,
            _ev("A", 2015, "revalued", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        b = lineage._to_lineage_event(
            store,
            _ev("B", 2021, "created", item_key="YBM4IZUS", p_id=686),
            mtime=0,
        )
        ev = _lineage_evidence([a, b], explicit=("A",))
        (out,) = lineage_sources(store, ev, max_items=1)
        assert out["label"] == a.label

    def test_element_diff_anchors_appended_after_events(self, store):
        e = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        diff = _diff("YBM4IZUS", 686, "CY2021 PFS final ¶686")
        ev = _lineage_evidence([e], element_diffs=[diff])
        out = lineage_sources(store, ev, max_items=8)
        assert [s["label"] for s in out] == [e.label, diff.label]
        assert out[1]["kind"] == "rule"
        assert out[1]["item_key"] == "YBM4IZUS"
        assert out[1]["p_id"] == "686"

    def test_element_diff_anchors_capped_at_two_most_shared_first(self, store):
        _insert_extra_anchor(store, "EX1", 2016, 10)
        _insert_extra_anchor(store, "EX2", 2017, 11)
        diffs = [
            _diff("DE2VH9PD", 1251, "CY2015 PFS final ¶1251"),  # most shared
            _diff("EX1", 10, "CY2016 PFS final ¶10"),
            _diff("EX2", 11, "CY2017 PFS final ¶11"),  # third — dropped by the cap
        ]
        ev = _lineage_evidence([], element_diffs=diffs)
        out = lineage_sources(store, ev, max_items=8)
        assert [s["label"] for s in out] == [diffs[0].label, diffs[1].label]

    def test_element_diffs_still_added_past_max_items_up_to_the_hard_cap(self, store):
        # 8 distinct event paragraphs (8 distinct codes, so Ruling B11's
        # 4-per-code priority cap doesn't itself trim the selection)
        # exactly fill max_items=8; two more distinct element-diff
        # paragraphs are still appended (the "add up to 2 even when
        # max_items has no room" rule), reaching the absolute hard cap
        # of 10, not 11.
        events = []
        for i in range(8):
            key = f"EV{i}"
            _insert_extra_anchor(store, key, 2000 + i, 100 + i)
            events.append(
                lineage._to_lineage_event(
                    store,
                    _ev(f"Z{i}", 2000 + i, "created", item_key=key, p_id=100 + i),
                    mtime=0,
                )
            )
        _insert_extra_anchor(store, "DX1", 2050, 500)
        _insert_extra_anchor(store, "DX2", 2051, 501)
        _insert_extra_anchor(store, "DX3", 2052, 502)
        diffs = [
            _diff("DX1", 500, "d1"),
            _diff("DX2", 501, "d2"),
            _diff("DX3", 502, "d3"),  # third diff — dropped by the element cap of 2
        ]
        ev = _lineage_evidence(events, explicit=("Z0",), element_diffs=diffs)
        out = lineage_sources(store, ev, max_items=8)
        assert len(out) == 10
        assert [s["label"] for s in out[8:]] == ["d1", "d2"]

    def test_element_diff_with_no_anchor_is_skipped(self, store):
        e = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        rvu_only_diff = ElementDiff(
            type="t",
            value="v",
            in_codes=("A",),
            not_in_codes=("B",),
            label="l",
            item_key="",
            p_id=0,
        )
        ev = _lineage_evidence([e], element_diffs=[rvu_only_diff])
        out = lineage_sources(store, ev, max_items=8)
        assert [s["label"] for s in out] == [e.label]

    def test_element_diff_deduped_against_an_event_paragraph(self, store):
        e = lineage._to_lineage_event(
            store,
            _ev("99490", 2015, "adopted_cpt", item_key="DE2VH9PD", p_id=1251),
            mtime=0,
        )
        same_paragraph_diff = _diff("DE2VH9PD", 1251, "some other label")
        ev = _lineage_evidence([e], element_diffs=[same_paragraph_diff])
        out = lineage_sources(store, ev, max_items=8)
        assert len(out) == 1
        assert out[0]["label"] == e.label

    def test_element_diff_anchor_failure_is_skipped_not_raised(self, store):
        """One element-diff anchor failing to build (``_labeled_rule_source``
        raising) must be logged and skipped, not propagate — mirrors the
        event loop's own never-raise contract just above."""
        diff = _diff("DE2VH9PD", 1251, "CY2015 PFS final ¶1251")
        ev = _lineage_evidence([], element_diffs=[diff])
        with patch(
            "llm.lineage._labeled_rule_source", side_effect=RuntimeError("boom")
        ):
            out = lineage_sources(store, ev, max_items=8)
        assert out == []


class TestReconcileLabels:
    """``_reconcile_labels`` — a retrieved/cited source for the same FR
    paragraph as a lineage row wins the slot, renamed to the lineage
    label; the lineage row is dropped so it isn't merged as a
    duplicate under a second label."""

    def test_renames_matching_source_and_drops_the_lineage_row(self):
        retrieved = {
            "label": "85 FR 84639 ¶12",
            "id": "85 FR 84639 ¶12",
            "kind": "rule",
            "item_key": "YBM4IZUS",
            "p_id": "1578",
        }
        lineage_row = {
            "label": "CY2021 PFS final ¶1578",
            "id": "CY2021 PFS final ¶1578",
            "kind": "rule",
            "item_key": "YBM4IZUS",
            "p_id": "1578",
        }
        sources, remaining = _reconcile_labels([retrieved], [lineage_row])
        assert sources[0]["label"] == "CY2021 PFS final ¶1578"
        assert sources[0]["id"] == "CY2021 PFS final ¶1578"
        assert remaining == []

    def test_no_match_keeps_both_untouched(self):
        retrieved = {
            "label": "85 FR 1 ¶9",
            "id": "85 FR 1 ¶9",
            "kind": "rule",
            "item_key": "OTHER",
            "p_id": "9",
        }
        lineage_row = {
            "label": "CY2021 PFS final ¶1578",
            "id": "CY2021 PFS final ¶1578",
            "kind": "rule",
            "item_key": "YBM4IZUS",
            "p_id": "1578",
        }
        sources, remaining = _reconcile_labels([retrieved], [lineage_row])
        assert sources[0]["label"] == "85 FR 1 ¶9"
        assert remaining == [lineage_row]

    def test_non_rule_sources_are_never_renamed(self):
        comment = {
            "label": "CMS-2026-2377-1 p.1",
            "id": "CMS-2026-2377-1 p.1",
            "kind": "comment",
            "item_key": "YBM4IZUS",
            "p_id": "",
        }
        lineage_row = {
            "label": "CY2021 PFS final ¶1578",
            "id": "CY2021 PFS final ¶1578",
            "kind": "rule",
            "item_key": "YBM4IZUS",
            "p_id": "1578",
        }
        sources, remaining = _reconcile_labels([comment], [lineage_row])
        assert sources[0]["label"] == "CMS-2026-2377-1 p.1"
        assert remaining == [lineage_row]


class TestLineageEvidenceLive:
    @pytest.mark.live
    def test_median_under_300ms_on_live_replica(self, restore_families):
        # I5: lineage_evidence -> evidence._connect -> refresh_from(con)
        # merges the live replica's derived pfs.code_family rows into the
        # module-global FAMILIES registry — restore_families snapshots and
        # restores it so this test doesn't leak thousands of derived
        # families into every test that runs after it in the session.
        from conf import ROOT

        replica = ROOT / "data" / "replica" / "aco.ro.duckdb"
        if not replica.exists():
            pytest.skip(f"no live replica at {replica}")
        cfg = replace(CFG, duckdb_replica=str(replica))
        lineage_evidence("history of CCM coding and payment", cfg)  # warm
        times_ms = []
        for _ in range(5):
            start = time.perf_counter()
            lineage_evidence("history of CCM coding and payment", cfg)
            times_ms.append((time.perf_counter() - start) * 1000)
        times_ms.sort()
        median = times_ms[len(times_ms) // 2]
        print(f"lineage_evidence median: {median:.2f} ms (all: {times_ms})")
        assert median < 300, f"{median:.2f} ms over {times_ms}"


class TestKnownCodeFilter:
    """FR page numbers leak into from/to lists ("(84 FR 62694 through
    62695)" → 62695); with the replica's known-code set they are dropped
    and an empty directional row disappears."""

    def test_junk_tokens_dropped_and_empty_directional_row_removed(self):
        known = frozenset({"G2058", "99439", "99487"})
        row = _ev("G2058", 2021, "replaced_by", to="84750 99439 62694", source="fr")
        ev = lineage._to_lineage_event(None, row, 1, known)
        assert ev is not None and ev.to_codes == ("99439",)
        junk = _ev("G2058", 2022, "replaced_by", to="84755 84756", source="fr")
        assert lineage._to_lineage_event(None, junk, 1, known) is None
        # non-directional rows keep going even with nothing left
        created = _ev("G2058", 2021, "created", to="84750", source="fr")
        assert lineage._to_lineage_event(None, created, 1, known).to_codes == ()

    def test_no_known_set_means_no_filtering(self):
        row = _ev("G2058", 2021, "replaced_by", to="84750 99439", source="fr")
        assert lineage._to_lineage_event(None, row, 1).to_codes == ("84750", "99439")

    def test_known_codes_reads_both_tables_and_tolerates_missing(self, con):
        c, _ = con
        c.execute("CREATE TABLE pfs.rvu (hcpcs VARCHAR, year INTEGER)")
        c.execute("INSERT INTO pfs.rvu VALUES ('99490', 2026)")
        known = lineage._known_codes(c, 1)
        assert "99490" in known
        import duckdb

        bare = duckdb.connect(":memory:")
        assert lineage._known_codes(bare, 2) == frozenset()

    def test_known_codes_logs_and_continues_on_a_real_error(self):
        """A missing table (I4) is expected and silent; any other
        failure is still swallowed (the known-code filter must never
        break the chat) but is worth a log line — this exercises that
        branch specifically, distinct from the missing-table path
        above."""

        class _FakeCur:
            def execute(self, sql):
                raise ValueError("boom")

        assert lineage._known_codes(_FakeCur(), 777_001) == frozenset()

    def test_priority_kinds_always_kept_non_priority_trimmed(self):
        priority = LineageEvent(
            "99490", 2021, "created", (), (), "[L]", "K", 1, 0, "", "fr", True, ""
        )
        noisy = [
            LineageEvent(
                "99490",
                2022 + i,
                "revalued",
                (),
                (),
                f"[L{i}]",
                "K",
                1,
                0,
                "",
                "rvu",
                True,
                "",
            )
            for i in range(5)
        ]
        ev = LineageEvidence(
            codes=("99490",),
            families=(),
            events=(priority, *noisy),
            element_diffs=(),
            guidance=(),
            max_prompt_rows=3,
        )
        rendered = ev.prompt_block()
        assert "created" in rendered
        assert rendered.count("revalued") == 2  # budget: 3 - 1 priority = 2 kept