"""Fetch + parse public CMS-HCC v28 model artifacts into data/reference/.

CMS renames the model-software files per release (the 2024 zip is named
differently from 2025), so this script takes the zip URL as an argument and
resolves the four files we need *inside* the archive by content-sniffing,
not by filename.

Usage:
    uv run python scripts/fetch_cms_hcc_v28.py hydrate \\
        --zip-url "https://www.cms.gov/files/zip/<file>.zip" \\
        --icd10-csv-url "https://www.cms.gov/files/<icd10-descriptions>.csv" \\
        --out data/reference/

Discover the correct URLs on:
    https://www.cms.gov/medicare/payment/medicare-advantage-rates-statistics/risk-adjustment/2024-model-software/icd-10-mappings
    https://www.cdc.gov/nchs/icd/icd-10-cm.htm

The script:
  1. Downloads the zip + optional ICD-10 description file.
  2. Verifies the SHA-256 (logs it; no expected value required).
  3. Unpacks into a temp dir and scans for the four artifacts:
       - a CSV / XLSX / SAS file whose columns resemble ICD-10 → HCC
       - the HCC label / description table
       - the hierarchy (trumping) table
       - the RAF coefficients table (Community Non-Dual Aged + Dual + Institutional)
  4. Normalizes each to the seed schema used by `app/pipeline/components/e_crosswalk.py`
     and writes the four CSVs into --out.
  5. Emits a `PROVENANCE.lock` next to the CSVs with source URLs, SHA-256 digests,
     timestamps, and file list.

This script is intentionally conservative: if it can't auto-detect a file, it
prints the extracted file listing and exits non-zero so Charles can point it
at the right table with a follow-up flag (--coefficients-name etc).

Implementation status: **partial**. Download + checksum + enumerate is wired;
the per-file normalizer is stubbed for the CMS-specific column names. Fill in
the normalizer once the zip is in hand and we know the exact headers.
"""

from __future__ import annotations

import csv
import hashlib
import io
import json
import re
import sys
import tempfile
import zipfile
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

import httpx
import typer
from rich.console import Console

app = typer.Typer(help=__doc__, no_args_is_help=True)
console = Console()

REF_DIR_DEFAULT = Path("data/reference")


@dataclass
class DownloadRecord:
    url: str
    local_path: Path
    size_bytes: int
    sha256: str


def _download(url: str, out_path: Path, timeout_s: int = 120) -> DownloadRecord:
    console.print(f"Downloading {url}")
    sha = hashlib.sha256()
    size = 0
    with httpx.stream("GET", url, follow_redirects=True, timeout=timeout_s) as resp:
        resp.raise_for_status()
        with out_path.open("wb") as f:
            for chunk in resp.iter_bytes(chunk_size=1 << 16):
                f.write(chunk)
                sha.update(chunk)
                size += len(chunk)
    record = DownloadRecord(
        url=url, local_path=out_path, size_bytes=size, sha256=sha.hexdigest()
    )
    console.print(
        f"  → {out_path} ({size / 1024:.1f} KiB, sha256={record.sha256[:12]}…)"
    )
    return record


def _enumerate_zip(zip_path: Path, extract_dir: Path) -> list[Path]:
    console.print(f"Extracting {zip_path}")
    with zipfile.ZipFile(zip_path) as zf:
        zf.extractall(extract_dir)
    files = sorted(p for p in extract_dir.rglob("*") if p.is_file())
    console.print(f"  Found {len(files)} files:")
    for p in files:
        console.print(f"    {p.relative_to(extract_dir)}  ({p.stat().st_size} bytes)")
    return files


# ---------------------------------------------------------------- sniffers

_ICD_COL_PATTERNS = [re.compile(r"^icd[\s_-]*10([\s_-]*cm)?$", re.I), re.compile(r"^dx[_\s]*code$", re.I)]
_HCC_COL_PATTERNS = [re.compile(r"^hcc(_?id)?$", re.I), re.compile(r"^cc$", re.I)]
_LABEL_COL_PATTERNS = [re.compile(r"label", re.I), re.compile(r"descr", re.I)]
_RANK_COL_PATTERNS = [re.compile(r"rank", re.I)]
_GROUP_COL_PATTERNS = [re.compile(r"hier", re.I), re.compile(r"group", re.I)]


def _sniff_table(path: Path) -> list[dict[str, str]] | None:
    """Load a CSV or XLSX-ish table as a list of dicts. Returns None for unknown formats."""
    suffix = path.suffix.lower()
    if suffix in {".csv", ".tsv", ".txt"}:
        delim = "\t" if suffix == ".tsv" else None
        with path.open(newline="", encoding="utf-8", errors="replace") as f:
            sample = f.read(2048)
            f.seek(0)
            if delim is None:
                try:
                    delim = csv.Sniffer().sniff(sample, delimiters=",\t|").delimiter
                except csv.Error:
                    delim = ","
            return list(csv.DictReader(f, delimiter=delim))
    if suffix in {".xlsx", ".xls"}:
        try:
            import pandas as pd

            df = pd.read_excel(path)
            return [{str(k): ("" if pd.isna(v) else str(v)) for k, v in row.items()} for row in df.to_dict("records")]
        except ImportError:
            console.print("[yellow]pandas/openpyxl not installed — skipping xlsx sniff[/yellow]")
            return None
    return None


def _classify(rows: list[dict[str, str]]) -> str | None:
    if not rows:
        return None
    headers = {h.lower(): h for h in rows[0].keys()}
    if any(pat.search(h) for h in headers for pat in _ICD_COL_PATTERNS) and any(
        pat.search(h) for h in headers for pat in _HCC_COL_PATTERNS
    ):
        return "crosswalk"
    if any(pat.search(h) for h in headers for pat in _HCC_COL_PATTERNS) and any(
        pat.search(h) for h in headers for pat in _RANK_COL_PATTERNS
    ):
        return "hierarchy"
    if any(pat.search(h) for h in headers for pat in _HCC_COL_PATTERNS) and any(
        pat.search(h) for h in headers for pat in _LABEL_COL_PATTERNS
    ):
        return "labels"
    if any("coeff" in h.lower() for h in headers):
        return "coefficients"
    return None


# ---------------------------------------------------------------- commands


@app.command()
def hydrate(
    zip_url: str = typer.Option(..., help="CMS v28 model zip URL."),
    icd10_csv_url: str | None = typer.Option(None, help="Optional CDC ICD-10-CM descriptions CSV URL."),
    out: Path = typer.Option(REF_DIR_DEFAULT, help="Output directory."),
    keep_raw: bool = typer.Option(False, help="Keep the downloaded zip + extracted files on disk."),
) -> None:
    """Download → extract → sniff → normalize → write four seed CSVs + PROVENANCE.lock."""
    out.mkdir(parents=True, exist_ok=True)

    with tempfile.TemporaryDirectory() as td:
        td_path = Path(td)
        zip_rec = _download(zip_url, td_path / "cms_hcc_v28.zip")

        icd_rec: DownloadRecord | None = None
        if icd10_csv_url:
            icd_rec = _download(icd10_csv_url, td_path / "icd10_cm_descriptions.csv")
            _copy_icd_descriptions(icd_rec.local_path, out / "icd10_cm_descriptions.csv")

        extracted = _enumerate_zip(zip_rec.local_path, td_path / "extracted")
        classified: dict[str, Path] = {}
        for p in extracted:
            rows = _sniff_table(p)
            kind = _classify(rows) if rows else None
            if kind and kind not in classified:
                classified[kind] = p
                console.print(f"  identified {p.name} → {kind}")

        missing = {"crosswalk", "labels", "hierarchy"} - set(classified)
        if missing:
            console.print(
                f"[red]Could not auto-classify:[/red] {sorted(missing)}"
            )
            console.print("Re-run with explicit --xxx-name flags (TODO: plumb them).")
            console.print(f"Extraction staged at: {td_path / 'extracted'}")
            raise typer.Exit(2)

        if keep_raw:
            keep_dir = out.parent / "_cms_raw"
            keep_dir.mkdir(parents=True, exist_ok=True)
            zip_rec.local_path.rename(keep_dir / zip_rec.local_path.name)

        # Normalizers are intentionally stubbed — they depend on the exact column
        # names CMS ships, which change year to year. Print a signal and exit
        # non-zero so Charles can inspect + add mapping rules in a follow-up.
        console.print(
            "[yellow]Column normalizers not yet wired.[/yellow] "
            "Inspect the classified files and extend this script with:\n"
            "  _normalize_crosswalk(...) → writes hcc_v28_crosswalk.csv\n"
            "  _normalize_labels(...)    → writes hcc_v28_labels.csv\n"
            "  _normalize_hierarchy(...) → writes hcc_v28_hierarchy.csv\n"
            "  _normalize_coefficients(...) → merges RAF weights into labels.csv"
        )
        console.print("Classified files:")
        for kind, path in classified.items():
            console.print(f"  {kind:12s} → {path}")

        lock = {
            "generated_at": datetime.now(timezone.utc).isoformat(),
            "zip": {"url": zip_rec.url, "sha256": zip_rec.sha256, "bytes": zip_rec.size_bytes},
            "icd10_csv": (
                {"url": icd_rec.url, "sha256": icd_rec.sha256, "bytes": icd_rec.size_bytes}
                if icd_rec
                else None
            ),
            "classified": {k: str(p.relative_to(td_path)) for k, p in classified.items()},
        }
        (out / "PROVENANCE.lock").write_text(json.dumps(lock, indent=2))
        console.print(f"[green]Wrote {out / 'PROVENANCE.lock'}[/green]")


def _copy_icd_descriptions(src: Path, dst: Path) -> None:
    """Copy the CDC ICD-10-CM descriptions into our schema.

    CDC ships two common formats: a fixed-width `.txt` and a CSV-ish table. We
    detect CSV first; if the file looks fixed-width (first line lacks commas
    but starts with alphanumeric code), fall back to a column-offset parse.
    """
    text = src.read_text(encoding="utf-8", errors="replace")
    first = text.splitlines()[0] if text else ""
    writer_rows: list[tuple[str, str]] = []

    if "," in first and first.lower().startswith(("icd", "code", '"icd', '"code')):
        reader = csv.DictReader(io.StringIO(text))
        for row in reader:
            code = (row.get("icd10_code") or row.get("code") or "").strip()
            desc = (row.get("description") or row.get("long_description") or row.get("desc") or "").strip()
            if code and desc:
                writer_rows.append((code, desc))
    else:
        # Fixed-width heuristic: code starts at col 0, description starts after the
        # first contiguous whitespace gap.
        for line in text.splitlines():
            m = re.match(r"^(\S+)\s+(.+)$", line)
            if m:
                writer_rows.append((m.group(1).strip(), m.group(2).strip()))

    with dst.open("w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["icd10_code", "description"])
        w.writerows(writer_rows)
    console.print(f"  wrote {dst} ({len(writer_rows)} rows)")


@app.command()
def inspect(
    zip_url: str = typer.Option(..., help="CMS v28 zip URL — just download + list contents."),
) -> None:
    """Download the zip and list its contents without touching data/reference/.

    Use this to discover the exact filenames before running `hydrate`.
    """
    with tempfile.TemporaryDirectory() as td:
        zp = Path(td) / "zip.zip"
        _download(zip_url, zp)
        extract = Path(td) / "extract"
        files = _enumerate_zip(zp, extract)
        console.print(f"\nTotal: {len(files)} files.")


if __name__ == "__main__":
    try:
        app()
    except httpx.HTTPError as e:
        console.print(f"[red]Network error: {e}[/red]")
        sys.exit(1)