stack/tests/pfs/test_commenters.py

"""pfs.commenters — commenter organization, type, specialty, campaigns (task 8)."""

from __future__ import annotations

from types import SimpleNamespace
from unittest.mock import MagicMock

import duckdb
import httpx
import pytest

from bib.item import Source
from bib.store import Store
from pfs.codetables import (
    CommenterRow,
    ensure_tables,
    read_commenters,
    write_commenters,
)
from pfs.commenters import (
    CATEGORY_MAP,
    IDENTITY_PROMPT_VERSION,
    ORG_TYPES,
    MirrorUnavailable,
    RegsFields,
    family_comment_items,
    form_groups,
    identify,
    item_openings,
    letterhead_identity,
    regs_fields,
    specialty_key,
    specialty_vocab,
)
from pfs.runlog import RunLog

HOME = SimpleNamespace(
    key="HOME-VISITS", codes=("99341", "99342"), pair_phrases=("home or residence",)
)


class FakeMirror:
    def __init__(self, attrs=None, per_id=None, fail=False):
        self.attrs = attrs or {}
        self.per_id = per_id or {}
        self.fail = fail

    def get_json(self, docket, cid):
        if self.fail:
            raise RuntimeError("404")
        return {"data": {"attributes": self.per_id.get(cid, self.attrs)}}


class FakeAsk:
    model = "qwen-test"

    def __init__(self, reply):
        self.reply = reply
        self.calls = []

    def __call__(self, system, user, schema):
        self.calls.append((system, user, schema))
        return self.reply


def engine_with(rows, openings=None, extra=()):
    """Fake pgvector engine that emulates the two chunk queries. *rows* are
    the family-matching ``(item_key, text)`` chunks (late in the comment);
    *openings* ``{item_key: [chunk texts]}`` are non-matching opening chunks of
    the inline comment; *extra* are further ``(item, section, seq, doc,
    matched)`` chunks. The query that binds ``keys`` is the openings query."""
    engine = MagicMock()
    conn = engine.begin.return_value.__enter__.return_value
    chunks = [(k, "Inline comment", 100, t, True) for k, t in rows]
    for k, parts in (openings or {}).items():
        chunks += [(k, "Inline comment", i, t, False) for i, t in enumerate(parts)]
    chunks += list(extra)

    def execute(sql, params):
        by_sec = {}
        for c in sorted(chunks, key=lambda c: c[2]):
            by_sec.setdefault((c[0], c[1]), []).append(c)
        if "keys" in params:
            scope = set(params["keys"])
        else:
            scope = {c[0] for c in chunks if c[4]}
        out = []
        for secs in by_sec.values():
            for rs, c in enumerate(secs, 1):
                matched = c[4] and "keys" not in params
                if c[0] in scope and (rs <= 6 or matched):
                    out.append((c[0], c[1], c[2], c[3], rs, matched))
        result = MagicMock()
        result.fetchall.return_value = sorted(out, key=lambda r: (r[0], r[2]))
        return result

    conn.execute.side_effect = execute
    return engine, conn


@pytest.fixture
def con():
    c = duckdb.connect(":memory:")
    ensure_tables(c)
    c.execute("CREATE SCHEMA terminology")
    c.execute(
        "CREATE TABLE terminology.cms_acceptable_provider_specialty_codes "
        "(specialty_code VARCHAR, specialty_description VARCHAR, "
        "pecos_specialty_description VARCHAR)"
    )
    c.executemany(
        "INSERT INTO terminology.cms_acceptable_provider_specialty_codes "
        "VALUES (?,?,?)",
        [
            ("08", "Family Practice", ""),
            ("11", "Internal Medicine", ""),
            ("01", "Allergy/Immunology", ""),
            ("x", None, ""),
        ],
    )
    yield c
    c.close()


@pytest.fixture
def store():
    s = Store(":memory:")
    for i, k in enumerate(["K1", "K2", "K3", "K4", "NOURL"]):
        url = (
            "https://example.org/nothing"
            if k == "NOURL"
            else f"https://www.regulations.gov/comment/CMS-X-{i}"
        )
        s.create(Source(key=k, title=f"T{k}", url=url))
    return s


def _reply(**kw):
    base = {
        "organization": "",
        "org_type": "unknown",
        "specialty": "NONE",
        "on_behalf_of_organization": False,
    }
    return {**base, **kw}


def _run(tmp_path, con, store, mirror, ask, rows, key="K1"):
    return identify(
        engine_with(rows)[0],
        store,
        con,
        HOME,
        docket="CMS-2022-0113",
        mirror=mirror,
        ask=ask,
        log=RunLog(tmp_path / "l.jsonl"),
    )


def test_submitter_name_is_the_comments_own_posted_name(tmp_path, con, store):
    mirror = FakeMirror(
        {
            "category": "Physician",
            "organization": "",
            "firstName": " Jane ",
            "lastName": "Doe",
        }
    )
    ask = FakeAsk(_reply(org_type="individual_clinician", specialty="FAMILY PRACTICE"))
    rows, _ = _run(tmp_path, con, store, mirror, ask, [("K1", "Sincerely, Jane Doe")])
    assert rows[0].submitter_name == "Jane Doe"
    assert (rows[0].org_type, rows[0].org_name) == ("individual_clinician", "")
    write_commenters(con, rows)
    assert read_commenters(con)[0].submitter_name == "Jane Doe"


def test_two_comments_never_swap_names(tmp_path, con, store):
    # the fake store gives K1 .../CMS-X-0 and K2 .../CMS-X-1
    mirror = FakeMirror(
        per_id={
            "CMS-X-0": {"category": "Individual", "firstName": "Ann", "lastName": "A"},
            "CMS-X-1": {"category": "Individual", "firstName": "Bob", "lastName": "B"},
        }
    )
    ask = FakeAsk(_reply(org_type="individual"))
    rows, _ = _run(
        tmp_path, con, store, mirror, ask, [("K1", "letter one"), ("K2", "letter two")]
    )
    assert {r.item_key: r.submitter_name for r in rows} == {
        "K1": "Ann A",
        "K2": "Bob B",
    }
    assert {r.item_key: r.comment_id for r in rows} == {
        "K1": "CMS-X-0",
        "K2": "CMS-X-1",
    }


def test_solo_practice_keeps_org_name(tmp_path, con, store):
    mirror = FakeMirror(
        {"category": "", "organization": "", "firstName": "Jane", "lastName": "Doe"}
    )
    ask = FakeAsk(
        _reply(
            organization="Doe Family Medicine PLLC",
            org_type="health_system",
            on_behalf_of_organization=True,
        )
    )
    rows, _ = _run(
        tmp_path,
        con,
        store,
        mirror,
        ask,
        [("K1", "Letter of Doe Family Medicine PLLC about visits")],
    )
    assert (rows[0].org_name, rows[0].org_type) == (
        "Doe Family Medicine PLLC",
        "health_system",
    )
    assert rows[0].submitter_name == "Jane Doe"
    assert rows[0].basis == "letterhead_model"


def test_regs_organization_is_kept_even_with_submitter_name():
    m = FakeMirror(
        {
            "category": "Individual",
            "organization": "Doe Clinic",
            "firstName": "Jane",
            "lastName": "Doe",
        }
    )
    assert regs_fields(m, "D", "c") == RegsFields(
        "Individual", "Doe Clinic", "Jane Doe"
    )


def test_missing_name_is_blank():
    for attrs in ({}, {"firstName": None, "lastName": ""}, {"firstName": "  "}):
        assert regs_fields(FakeMirror(attrs), "D", "c").submitter_name == ""
    assert (
        regs_fields(FakeMirror({"lastName": "Doe"}), "D", "c").submitter_name == "Doe"
    )


def test_regs_fields_fails_closed_on_unreadable_json():
    with pytest.raises(MirrorUnavailable):
        regs_fields(FakeMirror(fail=True), "D", "c")
    odd = MagicMock()
    odd.get_json.return_value = {"data": {"attributes": "x"}}
    with pytest.raises(MirrorUnavailable):
        regs_fields(odd, "D", "c")
    odd.get_json.return_value = None
    with pytest.raises(MirrorUnavailable):
        regs_fields(odd, "D", "c")
    # a one-letter name never matches ("A" in "A Plus Health")
    m = FakeMirror({"organization": "A Plus Health", "firstName": "A"})
    assert regs_fields(m, "D", "c").organization == "A Plus Health"


def test_org_must_appear_in_text():
    ask = FakeAsk(
        _reply(
            organization="AAFP",
            org_type="specialty_society",
            on_behalf_of_organization=True,
        )
    )
    out = letterhead_identity(
        "Dear CMS, we write about visits.", ask, specialties=("X",)
    )
    assert out["organization"] == ""
    assert out["org_type"] == "specialty_society"
    # present in text (case/whitespace-normalized) -> kept
    out = letterhead_identity("Dear CMS,  the  aafp  writes", ask, specialties=("X",))
    assert out["organization"] == "AAFP"
    # present only in the regulations.gov organization -> kept
    out = letterhead_identity("x", ask, specialties=("X",), reg_org="AAFP Inc")
    assert out["organization"] == "AAFP"


def test_org_blanked_when_not_on_behalf_or_person_type_or_after_1200():
    base = dict(organization="AAFP", org_type="specialty_society")
    t = "AAFP letter"
    for reply in (
        _reply(**base, on_behalf_of_organization=False),
        _reply(
            organization="AAFP", org_type="individual", on_behalf_of_organization=True
        ),
    ):
        assert (
            letterhead_identity(t, FakeAsk(reply), specialties=("X",))["organization"]
            == ""
        )
    late = "z" * 1200 + " AAFP"
    ask = FakeAsk(_reply(**base, on_behalf_of_organization=True))
    assert letterhead_identity(late, ask, specialties=("X",))["organization"] == ""
    assert len(ask.calls[0][1]) == 1200


def test_letterhead_normalizes_bad_enums_and_none():
    out = letterhead_identity(
        "t",
        FakeAsk(_reply(org_type="bogus", specialty="NOPE")),
        specialties=("FAMILY PRACTICE",),
    )
    assert (out["org_type"], out["specialty"]) == ("unknown", "NONE")
    assert letterhead_identity("t", FakeAsk(None), specialties=()) is None
    schema = FakeAsk(None)
    letterhead_identity("t", schema, specialties=("A", "B"))
    props = schema.calls[0][2]["properties"]
    assert props["specialty"]["enum"] == ["A", "B", "NONE"]
    assert tuple(props["org_type"]["enum"]) == ORG_TYPES


def test_category_map_precedence(tmp_path, con, store):
    mirror = FakeMirror({"category": "Insurance", "organization": "Acme Health Plan"})
    ask = FakeAsk(
        _reply(
            organization="Vendor Inc", org_type="vendor", on_behalf_of_organization=True
        )
    )
    rows, _ = identify(
        engine_with([("K1", "Vendor Inc writes")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=mirror,
        ask=ask,
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert (rows[0].org_type, rows[0].basis) == ("payer", "regs_gov_fields")
    assert rows[0].org_name == "Acme Health Plan"
    assert rows[0].reg_category == "Insurance"
    assert CATEGORY_MAP["Insurance"] == "payer"


def test_basis_both_and_unknown(tmp_path, con, store):
    mirror = FakeMirror({"category": "Insurance"})
    ask = FakeAsk(_reply(org_type="payer", specialty="MULTISPECIALTY"))
    rows, _ = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=mirror,
        ask=ask,
        log=RunLog(tmp_path / "a.jsonl"),
    )
    assert rows[0].basis == "both" and rows[0].specialty == "MULTISPECIALTY"
    # no model answer and no mapped category: unknown, regs basis, extractor regs
    rows, counts = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror({"category": "Weird"}),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "b.jsonl"),
    )
    assert (rows[0].org_type, rows[0].basis, rows[0].extractor) == (
        "unknown",
        "regs_gov_fields",
        "regs",
    )
    assert counts["failed"] == 1 and counts["pending"] == 0


def test_resume_limit_and_pending(tmp_path, con, store):
    rows_in = [("K1", "one"), ("K2", "two"), ("K3", "three")]
    ask = FakeAsk(_reply(org_type="other"))
    log = RunLog(tmp_path / "l.jsonl")
    seen = []
    rows, counts = identify(
        engine_with(rows_in)[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=ask,
        log=log,
        limit=2,
        progress=lambda u, s: seen.append((u, s)),
    )
    assert counts["done"] == 2 and counts["pending"] == 1 and len(ask.calls) == 2
    assert seen == [("K1", "done"), ("K2", "done")]
    assert rows[2].extractor == "regs"
    rows, counts = identify(
        engine_with(rows_in)[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=ask,
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert (counts["skipped"], counts["done"], counts["pending"]) == (2, 1, 0)
    assert rows[0].extractor == f"model:qwen-test@{IDENTITY_PROMPT_VERSION}"


def test_comment_id_from_store_url(tmp_path, con, store):
    rows, _ = identify(
        engine_with([("K2", "a"), ("NOURL", "b"), ("MISSING", "c")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "l.jsonl"),
    )
    by = {r.item_key: r.comment_id for r in rows}
    assert by == {"K2": "CMS-X-1", "NOURL": "", "MISSING": ""}


def test_form_groups_marks_campaign(tmp_path, con, store):
    base = (
        "I am writing to urge CMS to protect home visits for homebound older adults. "
        "These services keep patients out of the hospital and must be paid fairly. "
        "Please finalize the proposed rule without cuts to the home visit codes. "
    )
    texts = {
        "K1": base + "Alpha",
        "K2": base + "Beta",
        "K3": base + "Gamma",
        "K4": "short",
    }
    out = form_groups(texts)
    assert set(out) == {"K1", "K2", "K3"}
    assert {v for v in out.values()} == {("K1", 3)}
    # two copies are not a campaign
    assert form_groups({"K1": base, "K2": base}) == {}
    rows, _ = identify(
        engine_with(sorted(texts.items()))[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert [(r.form_group, r.form_size) for r in rows] == [
        ("K1", 3),
        ("K1", 3),
        ("K1", 3),
        ("", 0),
    ]


def test_specialty_key_normalizes():
    assert specialty_key("Allergy/Immunology") == "ALLERGY IMMUNOLOGY"
    assert specialty_key("Hospice & Palliative  Care") == "HOSPICE AND PALLIATIVE CARE"
    assert specialty_key(None) == ""


def test_specialty_vocab(con):
    assert specialty_vocab(con) == (
        "ALLERGY IMMUNOLOGY",
        "FAMILY PRACTICE",
        "INTERNAL MEDICINE",
        "MULTISPECIALTY",
    )


def test_family_comment_items_joins_chunks_and_caps():
    engine, conn = engine_with(
        [("B", "b1"), ("A", "a1"), ("C", "x" * 4000), ("D", None)],
        extra=[("A", "Inline comment", 101, "a2", True)],
    )
    out = family_comment_items(
        engine, ["99341"], "home-visits", ["Home (or) residence"], docket="D1"
    )
    assert out == [("A", "a1\na2"), ("B", "b1"), ("C", "x" * 3000), ("D", "")]
    sql, params = conn.execute.call_args.args
    assert "~*" in str(sql) and params["docket"] == "D1"
    assert params["family_key"] == "HOME-VISITS" and params["codes"] == ["99341"]
    assert params["phrase_re"] == r"Home\ \(or\)\ residence"
    engine, conn = engine_with([])
    family_comment_items(engine, ["99341"], "X", [], docket="D1")
    sql, params = conn.execute.call_args.args
    assert "~*" not in str(sql) and "phrase_re" not in params


def test_family_comment_items_window_is_configurable():
    engine, _ = engine_with([("A", "a1"), ("B", "b" * 9000)])
    out = family_comment_items(
        engine, ["99341"], "H", [], docket="D1", max_chars=6000, max_chunks=8
    )
    assert out == [("A", "a1"), ("B", "b" * 6000)]


def test_commenter_table_roundtrip(con):
    def row(k, d="D", org="AAFP"):
        return CommenterRow(
            k, d, "c", org, "specialty_society", "", "", 0, "both", "Business", "regs"
        )

    assert write_commenters(con, [row("K1"), row("K2")]) == 2
    assert write_commenters(con, [row("K1", org="AMA")]) == 1
    got = read_commenters(con)
    assert [(r.item_key, r.org_name) for r in got] == [("K1", "AMA"), ("K2", "AAFP")]
    assert [r.item_key for r in read_commenters(con, ["K2"])] == ["K2"]
    assert read_commenters(con, []) == []
    assert write_commenters(con, []) == 0
    # a docket rewrite drops items the rerun no longer selects
    assert write_commenters(con, [row("K3")], docket="D") == 1
    assert [r.item_key for r in read_commenters(con)] == ["K3"]
    assert write_commenters(con, [row("K4", d="E")], docket="E") == 1
    assert [r.item_key for r in read_commenters(con)] == ["K3", "K4"]


def test_category_map_values_are_org_types_and_ambiguous_unmapped():
    assert set(CATEGORY_MAP.values()) <= set(ORG_TYPES)
    assert CATEGORY_MAP["Physician - HC005"] == "individual_clinician"
    assert CATEGORY_MAP["Hospital - HPA35"] == "health_system"
    for ambiguous in (
        "Health Care Professional/Association - Physician",
        "Association - Other",
        "Health Care Industry - PI015",
    ):
        assert CATEGORY_MAP.get(ambiguous, "") == ""


def test_identity_input_is_the_opening_not_the_matching_chunk(tmp_path, con, store):
    opening = {
        "K1": ["## Inline comment", "tMED Coalition writes.", "second chunk"]
        + [f"filler {i}" for i in range(4)]
    }
    engine, _ = engine_with(
        [("K1", "late matching chunk about 99341")], openings=opening
    )
    ask = FakeAsk(_reply(org_type="other"))
    identify(
        engine,
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=ask,
        log=RunLog(tmp_path / "l.jsonl"),
    )
    sent = ask.calls[0][1]
    assert sent.startswith("tMED Coalition writes.\nsecond chunk\nfiller 0")
    assert "late matching" not in sent and "##" not in sent
    assert IDENTITY_PROMPT_VERSION == "identity-v5"


RULE_HDR = (
    "## attachment_1.pdf DEPARTMENT OF HEALTH AND HUMAN SERVICES Centers for "
    "Medicare & Medicaid Services 42 CFR Parts 405, 410, 411"
)


def rule_upload(key, n=20, mention="99341 home visits"):
    """An attachment that is a copy of the rule: the FR header, then pages
    naming the code."""
    out = [(key, "attachment_1.pdf", 0, "## attachment_1.pdf", False)]
    out.append((key, "attachment_1.pdf", 1, RULE_HDR, False))
    out += [
        (key, "attachment_1.pdf", 2 + i, f"rule page {i} {mention}", True)
        for i in range(n)
    ]
    return out


def test_item_openings_headers_cap_and_rule_upload_left_out():
    engine, conn = engine_with(
        [],
        openings={"A": ["## Inline comment", "x" * 4000], "B": ["a letter"]},
        extra=rule_upload("B") + rule_upload("C"),
    )
    out = item_openings(engine, ["A", "B", "C", "Z"], docket="D1")
    assert out["A"].text == "x" * 3000 and not out["A"].rule_copy
    assert out["B"].text == "a letter" and out["B"].rule_copy
    assert "C" not in out and "Z" not in out  # only the rule: nothing of its own
    params = conn.execute.call_args.args[1]
    assert params["keys"] == ["A", "B", "C", "Z"] and params["docket"] == "D1"


def test_rule_copy_item_whose_inline_text_does_not_match_is_dropped(
    tmp_path, con, store
):
    engine, _ = engine_with(
        [("K1", "a letter naming 99341")],
        openings={"K2": ["A psychologists' form letter about medication."]},
        extra=rule_upload("K2"),
    )
    rows, counts = identify(
        engine,
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=FakeAsk(_reply()),
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert [r.item_key for r in rows] == ["K1"]
    assert family_comment_items(engine, ["99341"], "H", [], docket="D") == [
        ("K1", "a letter naming 99341")
    ]


def test_rule_copy_item_whose_inline_text_names_the_code_is_kept_without_the_rule(
    tmp_path, con, store
):
    engine, _ = engine_with(
        [("K2", "We support 99341 payment.")],
        openings={"K2": ["tMED Coalition writes."]},
        extra=rule_upload("K2"),
    )
    ask = FakeAsk(_reply(org_type="other"))
    rows, _ = identify(
        engine,
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=ask,
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert [r.item_key for r in rows] == ["K2"]
    sent = ask.calls[0][1]
    assert "tMED Coalition writes." in sent and "DEPARTMENT" not in sent
    assert "rule page" not in sent
    # selection text carries only the inline match, never a rule page
    assert family_comment_items(engine, ["99341"], "H", [], docket="D") == [
        ("K2", "We support 99341 payment.")
    ]


def test_header_must_be_in_the_attachment_opening_not_a_late_citation():
    late = [("A", "attachment_1.pdf", i, "page", False) for i in range(4)]
    late.append(("A", "attachment_1.pdf", 4, RULE_HDR, False))
    engine, _ = engine_with([], extra=late)
    assert item_openings(engine, ["A"], docket="D")["A"].rule_copy is False
    engine, _ = engine_with([], extra=[("B", "x", 0, "no section header here", False)])
    assert item_openings(engine, ["B"], docket="D")["B"].rule_copy is False


def test_form_groups_use_openings(tmp_path, con, store):
    letter = (
        "I am writing to urge CMS to protect home visits for homebound older adults. "
        "These services keep patients out of the hospital and must be paid fairly. "
        "Please finalize the proposed rule without cuts to the home visit codes. "
    )
    keys = ["K1", "K2", "K3"]
    # the matching excerpts differ; the openings are one form letter
    selector = [(k, f"different excerpt {k} " * 20) for k in keys]
    openings = {k: [letter + k] for k in keys}
    rows, _ = identify(
        engine_with(selector, openings=openings)[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "a.jsonl"),
    )
    assert [(r.form_group, r.form_size) for r in rows] == [("K1", 3)] * 3


def _status_error(code):
    req = httpx.Request("GET", "https://x")
    return httpx.HTTPStatusError(
        "e", request=req, response=httpx.Response(code, request=req)
    )


class SeqMirror:
    """Raises each queued exception in turn, then serves *attrs*."""

    def __init__(self, errors, attrs=None):
        self.errors = list(errors)
        self.attrs = attrs or {"category": "Individual"}
        self.calls = 0

    def get_json(self, docket, cid):
        self.calls += 1
        if self.errors:
            raise self.errors.pop(0)
        return {"data": {"attributes": self.attrs}}


@pytest.fixture(autouse=True)
def no_sleep(monkeypatch):
    monkeypatch.setattr("pfs.commenters.time.sleep", lambda s: None)


def test_mirror_retries_network_errors_and_5xx_twice():
    for errs in (
        [httpx.ConnectError("x"), httpx.ReadTimeout("y")],
        [_status_error(503), _status_error(500)],
    ):
        m = SeqMirror(errs)
        assert regs_fields(m, "D", "c").category == "Individual" and m.calls == 3
    m = SeqMirror([httpx.ConnectError("x")] * 3)
    with pytest.raises(MirrorUnavailable):
        regs_fields(m, "D", "c")
    assert m.calls == 3
    # a 404 or a parse problem is not retried
    m = SeqMirror([_status_error(404)])
    with pytest.raises(MirrorUnavailable):
        regs_fields(m, "D", "c")
    assert m.calls == 1


def test_blocked_item_becomes_unknown_row_without_model_call_then_is_replaced(
    tmp_path, con, store
):
    path = tmp_path / "l.jsonl"
    ask = FakeAsk(_reply(org_type="payer"))
    kw = dict(docket="D", ask=ask, log=RunLog(path))
    rows, counts = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        mirror=FakeMirror(fail=True),
        **kw,
    )
    assert counts["blocked"] == 1 and counts["pending"] == 0 and ask.calls == []
    r = rows[0]
    assert (r.item_key, r.org_name, r.org_type, r.specialty) == (
        "K1",
        "",
        "unknown",
        "",
    )
    assert r.extractor == "blocked:mirror" and not path.exists()
    write_commenters(con, rows)
    # the mirror reads on a later run: the row is replaced by the real one
    rows, counts = identify(
        engine_with([("K1", "a")])[0], store, con, HOME, mirror=FakeMirror(), **kw
    )
    assert counts["blocked"] == 0 and rows[0].org_type == "payer"
    write_commenters(con, rows)
    assert [r.org_type for r in read_commenters(con)] == ["payer"]


def test_retry_failed_reasks_failed_entries(tmp_path, con, store):
    log = RunLog(tmp_path / "l.jsonl")
    kw = dict(docket="D", mirror=FakeMirror(), log=log)
    ask = FakeAsk(None)
    _, counts = identify(engine_with([("K1", "a")])[0], store, con, HOME, ask=ask, **kw)
    assert counts["failed"] == 1
    _, counts = identify(engine_with([("K1", "a")])[0], store, con, HOME, ask=ask, **kw)
    assert counts["failed_pending"] == 1 and len(ask.calls) == 1
    ask.reply = _reply(org_type="payer")
    rows, counts = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        ask=ask,
        retry_failed=True,
        **kw,
    )
    assert counts["done"] == 1 and rows[0].org_type == "payer"


def test_government_and_other_categories_defer_to_the_model():
    for cat in CATEGORY_MAP:
        assert not cat.startswith(
            ("Government", "Federal Government", "Other", "Academic", "Attorney")
        )
    assert "Government - Federal" not in CATEGORY_MAP
    assert "other" not in CATEGORY_MAP.values()


def test_prompt_defines_vendor_and_sends_trade_associations_to_other():
    ask = FakeAsk(_reply())
    letterhead_identity("t", ask, specialties=())
    system = ask.calls[0][0]
    assert (
        "vendor (only a for-profit company" in system
        and "never an association" in system
    )
    assert "other (including trade associations, coalitions" in system


def test_family_rewrite_leaves_other_families_rows_alone(con):
    def row(k, fam, org="AAFP"):
        return CommenterRow(
            k, "D", "c", org, "specialty_society", "", "", 0, "both", "", "regs", fam
        )

    write_commenters(con, [row("K1", "A"), row("K2", "A")], docket="D", family="A")
    write_commenters(con, [row("K1", "B"), row("K3", "B")], docket="D", family="B")
    assert [(r.item_key, r.family) for r in read_commenters(con)] == [
        ("K1", "A"),
        ("K1", "B"),
        ("K2", "A"),
        ("K3", "B"),
    ]
    # rewriting A for D drops A's stale K2 but keeps all of B's
    write_commenters(con, [row("K1", "A", org="AMA")], docket="D", family="A")
    got = read_commenters(con)
    assert [(r.item_key, r.family, r.org_name) for r in got] == [
        ("K1", "A", "AMA"),
        ("K1", "B", "AAFP"),
        ("K3", "B", "AAFP"),
    ]
    assert [r.item_key for r in read_commenters(con, family="B")] == ["K1", "K3"]
    assert [r.item_key for r in read_commenters(con, ["K1"], family="A")] == ["K1"]


def test_legacy_commenter_table_is_migrated_to_home_visits():
    c = duckdb.connect(":memory:")
    c.execute("CREATE SCHEMA pfs")
    c.execute(
        "CREATE TABLE pfs.commenter (item_key VARCHAR, docket VARCHAR, comment_id VARCHAR, "
        "org_name VARCHAR, org_type VARCHAR, specialty VARCHAR, form_group VARCHAR, "
        "form_size INTEGER, basis VARCHAR, reg_category VARCHAR, extractor VARCHAR)"
    )
    c.execute(
        "INSERT INTO pfs.commenter VALUES ('K1','D','c','','unknown','','',0,'','','regs')"
    )
    ensure_tables(c)
    assert [r.family for r in read_commenters(c)] == ["HOME-VISITS"]
    # the name column is added too, blank for rows written before it existed
    assert [r.submitter_name for r in read_commenters(c)] == [""]


def test_identify_sets_family(tmp_path, con, store):
    rows, _ = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "l.jsonl"),
    )
    assert rows[0].family == "HOME-VISITS"
    rows, _ = identify(
        engine_with([("K1", "a")])[0],
        store,
        con,
        HOME,
        docket="D",
        mirror=FakeMirror(fail=True),
        ask=FakeAsk(None),
        log=RunLog(tmp_path / "m.jsonl"),
    )
    assert rows[0].extractor == "blocked:mirror" and rows[0].family == "HOME-VISITS"