stack/dev/scripts/add_zipcode_to_zotero.py

"""Add ZIP code carrier locality files to Zotero library.

Downloads from CMS, extracts, and creates Zotero items with proper
tags and storage-linked attachments.

Usage:
    uv run python dev/scripts/add_zipcode_to_zotero.py
"""

import subprocess
import zipfile
from pathlib import Path

from conf import path as _conf_path
from zot.db import TYPE_MAP, Db, generate_key, now_iso

ZOTERO_DB = str(_conf_path("db.zotero"))
ZOTERO_STORAGE = str(_conf_path("storage.zotero"))

ZIPCODE_FILES = {
    2016: {
        "zip": "/tmp/pfs_zipcode/zipcode_2016.zip",
        "url": "https://www.cms.gov/medicare/medicare-fee-for-service-payment/prospmedicarefeesvcpmtgen/downloads/2016-yearend-zipcodefile.zip",
        "title": "CY 2016 PFS Zip Carrier Locality — 2016 Year End",
    },
    2017: {
        "zip": "/tmp/pfs_zipcode/zipcode_2017.zip",
        "url": "https://www.cms.gov/medicare/medicare-fee-for-service-payment/prospmedicarefeesvcpmtgen/downloads/end-of-year-zip-code.zip",
        "title": "CY 2017 PFS Zip Carrier Locality — 2017 Year End",
    },
    2018: {
        "zip": "/tmp/pfs_zipcode/zipcode_2018.zip",
        "url": "https://www.cms.gov/medicare/medicare-fee-for-service-payment/prospmedicarefeesvcpmtgen/downloads/2018-yearend-zipcodefile.zip",
        "title": "CY 2018 PFS Zip Carrier Locality — 2018 Year End",
    },
    2019: {
        "zip": "/tmp/pfs_zipcode/zipcode_2019.zip",
        "url": "https://www.cms.gov/files/zip/2019-end-year-zip-code-file.zip",
        "title": "CY 2019 PFS Zip Carrier Locality — 2019 Year End",
    },
    2020: {
        "zip": "/tmp/pfs_zipcode/zipcode_2020.zip",
        "url": "https://www.cms.gov/files/zip/2020-end-year-zip-code-file.zip",
        "title": "CY 2020 PFS Zip Carrier Locality — 2020 Year End",
    },
    2021: {
        "zip": "/tmp/pfs_zipcode/zipcode_2021.zip",
        "url": "https://www.cms.gov/files/zip/2021-end-year-zip-code-file-revised-05/27/2022.zip",
        "title": "CY 2021 PFS Zip Carrier Locality — 2021 Year End",
    },
    2022: {
        "zip": "/tmp/pfs_zipcode/zipcode_2022.zip",
        "url": "https://www.cms.gov/files/zip/2022-end-year-zip-code-file.zip",
        "title": "CY 2022 PFS Zip Carrier Locality — 2022 Year End",
    },
    2023: {
        "zip": "/tmp/pfs_zipcode/zipcode_2023.zip",
        "url": "https://www.cms.gov/files/zip/2023-end-year-zip-code-file.zip",
        "title": "CY 2023 PFS Zip Carrier Locality — 2023 Year End",
    },
    2024: {
        "zip": "/tmp/pfs_zipcode/zipcode_2024.zip",
        "url": "https://www.cms.gov/files/zip/2024-end-year-zip-code-file.zip",
        "title": "CY 2024 PFS Zip Carrier Locality — 2024 Year End",
    },
    2025: {
        "zip": "/tmp/pfs_zipcode/zipcode_2025.zip",
        "url": "https://www.cms.gov/files/zip/2025-end-year-zip-code-file.zip",
        "title": "CY 2025 PFS Zip Carrier Locality — 2025 Year End",
    },
    2026: {
        "zip": "/tmp/pfs_zipcode/zipcode_2026.zip",
        "url": "https://www.cms.gov/files/zip/zip-code-carrier-locality-file-revised-11-18-2025.zip",
        "title": "CY 2026 PFS Zip Carrier Locality — January 2026",
    },
}


def add_zipcode_year(db: Db, year: int, info: dict) -> None:
    now = now_iso()

    # Extract ZIP
    extract_dir = Path(f"/tmp/pfs_zipcode/ex_{year}")
    extract_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(info["zip"]) as zf:
        zf.extractall(extract_dir)

    # Create parent webpage item
    parent_id = db.create_item(TYPE_MAP["webpage"], now=now)
    db.set_fields(
        parent_id,
        {
            "title": info["title"],
            "date": f"{year}-01-01",
            "url": info["url"],
            "accessDate": now,
            "websiteType": "Government Data Portal",
            "websiteTitle": "Centers for Medicare & Medicaid Services",
        },
    )

    # Tags: module:pfs + year:YYYY
    db.sync_tags(parent_id, ["module:pfs", f"year:{year}"])

    print(f"Created parent item for {info['title']}")

    # Add attachments for .txt and .xlsx files
    att_count = 0
    for filepath in sorted(extract_dir.iterdir()):
        ext = filepath.suffix.upper()
        if ext not in (".TXT", ".XLSX"):
            continue

        att_key = generate_key()
        storage_dir = Path(ZOTERO_STORAGE) / att_key
        subprocess.run(["sudo", "mkdir", "-p", str(storage_dir)], check=True)
        dest = storage_dir / filepath.name
        subprocess.run(["sudo", "cp", str(filepath), str(dest)], check=True)
        subprocess.run(
            ["sudo", "chown", "-R", "100999:100999", str(storage_dir)],
            check=True,
        )

        ct_map = {
            ".TXT": "text/plain",
            ".XLSX": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
        }
        content_type = ct_map.get(ext, "application/octet-stream")

        att_id = db.add_attachment(
            parent_id,
            key=att_key,
            content_type=content_type,
            path=f"storage:{filepath.name}",
        )
        db.set_field(att_id, "title", filepath.name)
        att_count += 1

    print(f"  Added {att_count} attachments for year {year}")


def main() -> None:
    for year, info in sorted(ZIPCODE_FILES.items()):
        zp = Path(info["zip"])
        if not zp.exists():
            print(f"ERROR: {zp} not found — download first")
            return

    with Db(ZOTERO_DB) as db:
        for year, info in sorted(ZIPCODE_FILES.items()):
            add_zipcode_year(db, year, info)
        db.commit()
        print("\nDone. Committed to Zotero database.")


if __name__ == "__main__":
    main()