stack/dev/scripts/fetch_nprm_addenda.py

"""Fetch + register PFS NPRM (proposed rule) addenda, by rule year.

Downloads a CY{year} PFS Proposed Rule "Addenda" ZIP, which contains
Addendum B (proposed RVUs by HCPCS), and optionally registers Addendum B
in bib against that year's NPRM rule item, tagged ``sup:{year}_PFS_NPRM``
so the loader can discover the *proposed* Addendum B without touching
final-rule files (which carry ``sup:{year}_PFS_FR`` — the two namespaces
must never collide).

Usage:
    uv run python dev/scripts/fetch_nprm_addenda.py --dry-run
    uv run python dev/scripts/fetch_nprm_addenda.py --year 2027
    uv run python dev/scripts/fetch_nprm_addenda.py --year 2027 --register
"""

from __future__ import annotations

import argparse
import re
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import httpx

ROOT = Path(__file__).resolve().parents[2]

USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36"
)

ADDENDUM_B_RE = re.compile(r"(?i)addendum[_ ]?b")
ADDENDUM_E_RE = re.compile(r"(?i)addendum[_ ]?e")


@dataclass(frozen=True)
class YearConfig:
    """Per-rule-year config: where to fetch the addenda ZIP and how to
    register Addendum B against the corresponding bib rule item."""

    year: int
    detail_page_url: str
    addenda_url: str
    dest_dir: Path
    addenda_zip_name: str
    rule_item_key: str
    addendum_b_title: str
    addendum_e_title: str
    register_tags: list[str]


CONFIGS: dict[int, YearConfig] = {
    2026: YearConfig(
        year=2026,
        detail_page_url=(
            "https://www.cms.gov/medicare/payment/fee-schedules/physician/"
            "federal-regulation-notices/cms-1832-p"
        ),
        # CY2026 PFS Proposed Rule (CMS-1832-P) Addenda ZIP — resolved from
        # the CMS-1832-P detail page above (the canonical
        # ".../pfs-federal-regulation-notices/cms-1832-p" URL 404s; this is
        # the current path found via the federal-regulation-notices
        # listing). Download link text on that page: "CY 2026 PFS Proposed
        # Rule Addenda - Updated 07/29/2025" (the URL path segment
        # literally contains "07/29/2025" for the update-date suffix —
        # that's CMS's own href, not a typo here).
        addenda_url=(
            "https://www.cms.gov/files/zip/"
            "cy-2026-pfs-proposed-rule-addenda-updated-07/29/2025.zip"
        ),
        dest_dir=ROOT / "data" / "cms" / "pfs_nprm" / "2026",
        addenda_zip_name="cy2026_pfs_nprm_addenda.zip",
        rule_item_key="2KVJ2HKX",
        addendum_b_title="CY2026 PFS NPRM Addendum B (proposed RVUs)",
        addendum_e_title="CY2026 PFS NPRM Addendum E (proposed GPCIs)",
        register_tags=["sup:2026_PFS_NPRM", "module:pfs", "file:rvu", "year:2026"],
    ),
    2027: YearConfig(
        year=2027,
        detail_page_url=(
            "https://www.cms.gov/medicare/payment/fee-schedules/physician/"
            "federal-regulation-notices/cms-1848-p"
        ),
        # CY2027 PFS Proposed Rule (CMS-1848-P) Addenda ZIP — resolved from
        # the CMS-1848-P detail page above (fallback if that 404s: the
        # federal-regulation-notices index page). Download link text on
        # that page: "CY 2027 PFS Proposed Rule Addenda (Updated
        # 07/21/2026)".
        addenda_url=(
            "https://www.cms.gov/files/zip/"
            "cy-2027-pfs-proposed-rule-addenda-updated-07-21-2026.zip"
        ),
        dest_dir=ROOT / "data" / "cms" / "pfs_nprm" / "2027",
        addenda_zip_name="cy2027_pfs_nprm_addenda.zip",
        rule_item_key="5ITGVDJV",
        addendum_b_title="CY2027 PFS NPRM Addendum B (proposed RVUs)",
        addendum_e_title="CY2027 PFS NPRM Addendum E (proposed GPCIs)",
        register_tags=["sup:2027_PFS_NPRM", "module:pfs", "file:rvu", "year:2027"],
    ),
}


def download_zip(url: str, dest: Path, *, retries: int = 3) -> bool:
    """Download the addenda ZIP if it doesn't already exist."""
    import time

    if dest.exists():
        print(f"  EXISTS {dest.name}")
        return True
    print(f"  GET    {dest.name} ...", end="", flush=True)
    for attempt in range(1, retries + 1):
        try:
            resp = httpx.get(
                url,
                headers={"User-Agent": USER_AGENT},
                timeout=120,
                follow_redirects=True,
            )
            if resp.status_code == 200:
                dest.write_bytes(resp.content)
                print(f" {len(resp.content) / 1024 / 1024:.1f} MB")
                return True
            if resp.status_code < 500 or attempt == retries:
                print(f" HTTP {resp.status_code}")
                return False
        except Exception as e:
            if attempt == retries:
                print(f" ERROR: {e}")
                return False
        time.sleep(2 * attempt)
    return False


def extract_zip(zip_path: Path, dest_dir: Path) -> list[Path]:
    """Unzip into dest_dir. Idempotent — extraction overwrites in place."""
    with zipfile.ZipFile(zip_path) as zf:
        names = [n for n in zf.namelist() if not n.endswith("/")]
        zf.extractall(dest_dir)
    return sorted(dest_dir / n for n in names)


def find_addendum(files: list[Path], pattern: re.Pattern[str]) -> Path | None:
    """Find the extracted file whose name matches ``pattern``.

    Excludes macOS ``__MACOSX/._*`` resource-fork junk that CMS's zips
    sometimes carry alongside the real file — those AppleDouble files
    match the same name pattern and suffix, so they'd otherwise be
    eligible candidates too.
    """
    candidates = [
        f
        for f in files
        if pattern.search(f.name)
        and f.suffix.lower() in (".xlsx", ".csv")
        and not f.name.startswith("._")
        and "__MACOSX" not in f.parts
    ]
    if not candidates:
        return None
    # Prefer .xlsx over .csv if both are present.
    candidates.sort(key=lambda f: f.suffix.lower() != ".xlsx")
    return candidates[0]


def find_addendum_b(files: list[Path]) -> Path | None:
    """Backward-compatible Addendum B finder."""
    return find_addendum(files, ADDENDUM_B_RE)


def preview_addendum_b(path: Path) -> None:
    """Print the header row (and a couple of data rows) for verification."""
    import polars as pl

    if path.suffix.lower() == ".xlsx":
        df = pl.read_excel(path, read_options={"header_row": None})
    else:
        df = pl.read_csv(path, n_rows=5, has_header=False)
    print(f"  columns/first rows of {path.name}:")
    with pl.Config(tbl_cols=-1, tbl_width_chars=200):
        print(df.head(5))


def _existing_attachment(store: Any, item_key: str, filename: str) -> str | None:
    """Return the attachment key on item_key whose filename matches, if any."""
    con = store._con()  # noqa: SLF001 — same pattern as bib/oig.py, bib/iom.py
    row = con.execute(
        """SELECT a.key FROM attachments a
           JOIN items i ON a.item_id = i.id
           WHERE i.key = ? AND a.filename = ?""",
        (item_key, filename),
    ).fetchone()
    return row["key"] if row else None


def _attach_guarded(store: Any, item_key: str, path: Path, title: str) -> None:
    """Attach ``path`` to ``item_key`` under ``title``, unless an
    attachment with that title already exists (attach_file has no dedup
    of its own, so every --register re-run would otherwise duplicate)."""
    existing_key = _existing_attachment(store, item_key, title)
    if existing_key is not None:
        print(
            f"SKIP attach — {item_key} already has an attachment "
            f"titled {title!r} (key {existing_key})"
        )
        return
    att_key = store.attach_file(item_key, path, title=title)
    con = store._con()  # noqa: SLF001 — same pattern as bib/oig.py, bib/iom.py
    row = con.execute(
        "SELECT storage_path FROM attachments WHERE key = ?", (att_key,)
    ).fetchone()
    storage_path = row["storage_path"] if row else "<unknown>"
    print(f"Attached {path.name} (key {att_key}) -> {storage_path}")


def register(
    cfg: YearConfig, addendum_b_path: Path, addendum_e_path: Path | None = None
) -> None:
    """Attach Addendum B (and E when present) to the NPRM rule item and
    tag it.

    Idempotent: attaches are title-guarded (see ``_attach_guarded``) and
    the tags re-apply harmlessly via add_tag's INSERT OR IGNORE.
    """
    from conf import connect

    store = connect.bib()

    _attach_guarded(store, cfg.rule_item_key, addendum_b_path, cfg.addendum_b_title)
    if addendum_e_path is not None:
        _attach_guarded(store, cfg.rule_item_key, addendum_e_path, cfg.addendum_e_title)

    for tag in cfg.register_tags:
        store.add_tag(cfg.rule_item_key, tag)
    print(f"Tagged {cfg.rule_item_key} with: {', '.join(cfg.register_tags)}")
    item = store.get(cfg.rule_item_key)
    print(f"Current tags on {cfg.rule_item_key}: {item.tags}")


def main() -> None:
    years = sorted(CONFIGS)
    parser = argparse.ArgumentParser(
        description=(
            "Fetch + register PFS NPRM addenda for a given rule year "
            f"(available: {', '.join(str(y) for y in years)})."
        )
    )
    parser.add_argument(
        "--year",
        type=int,
        choices=years,
        default=2026,
        help=(
            "Rule year to fetch — one of "
            f"{', '.join(str(y) for y in years)} (default: 2026)."
        ),
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print what would be downloaded/extracted without doing it.",
    )
    parser.add_argument(
        "--register",
        action="store_true",
        help="After fetching, attach Addendum B to the rule item and apply "
        "its tags for the selected --year.",
    )
    args = parser.parse_args()
    cfg = CONFIGS[args.year]

    zip_dest = cfg.dest_dir / cfg.addenda_zip_name

    if args.dry_run:
        print(f"Would GET  {cfg.addenda_url}")
        print(f"Would save {zip_dest}")
        print(f"Would extract into {cfg.dest_dir}")
        if args.register:
            print(
                f"Would attach Addendum B to {cfg.rule_item_key} "
                f"and tag with: {', '.join(cfg.register_tags)}"
            )
        return

    cfg.dest_dir.mkdir(parents=True, exist_ok=True)

    if not download_zip(cfg.addenda_url, zip_dest):
        raise SystemExit(f"Failed to download {cfg.addenda_url}")

    files = extract_zip(zip_dest, cfg.dest_dir)
    print(f"\nExtracted {len(files)} files to {cfg.dest_dir}:")
    for f in files:
        size_kb = f.stat().st_size / 1024
        print(f"  {f.relative_to(cfg.dest_dir)} ({size_kb:.0f} KB)")

    addendum_b = find_addendum_b(files)
    if addendum_b is None:
        raise SystemExit(
            "Could not find an Addendum B file (pattern "
            f"{ADDENDUM_B_RE.pattern!r}) among extracted files."
        )
    print(f"\nAddendum B: {addendum_b.relative_to(cfg.dest_dir)}")
    preview_addendum_b(addendum_b)

    # Addendum E (proposed GPCIs) is optional-but-expected: every NPRM
    # addenda ZIP to date has carried one, but its absence shouldn't
    # block the Addendum B registration this script exists for.
    addendum_e = find_addendum(files, ADDENDUM_E_RE)
    if addendum_e is None:
        print(f"\nNo Addendum E file matched {ADDENDUM_E_RE.pattern!r} — skipping.")
    else:
        print(f"\nAddendum E: {addendum_e.relative_to(cfg.dest_dir)}")

    if args.register:
        print()
        register(cfg, addendum_b, addendum_e)


if __name__ == "__main__":
    main()