#!/usr/bin/env python3
"""Download and process the CMS ICD-10-CM 2026 data into a compressed CSV.

This script:
1. Downloads the official CMS ICD-10-CM 2026 XML zip
2. Parses the tabular XML to extract code hierarchy (chapters, sections,
   categories, and individual codes with descriptions)
3. Writes a trimmed, gzip-compressed CSV for bundling with the package

Data source: https://www.cms.gov/medicare/coding-billing/icd-10-codes
Inspired by: https://github.com/rjake/ICD10-hierarchy

Run from the project root:
    python scripts/build_data.py
"""

from __future__ import annotations

import csv
import gzip
import io
import re
import zipfile
from pathlib import Path
from urllib.request import urlretrieve
from xml.etree import ElementTree as ET

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
CMS_ZIP_URL = "https://www.cms.gov/files/zip/2026-code-tables-tabular-and-index.zip"
XML_INNER_PATH = "Table and Index/icd10cm_tabular_2026.xml"
OUTPUT_DIR = Path(__file__).resolve().parent.parent / "src" / "icdcodex" / "_data"
OUTPUT_FILE = OUTPUT_DIR / "icd10_2026.csv.gz"
TEMP_ZIP = Path(__file__).resolve().parent / "_cms_2026.zip"

# Output columns
COLUMNS = [
    "icd10_code",
    "description",
    "chapter_desc",
    "section_desc",
    "category_desc",
]


def download_zip() -> Path:
    """Download the CMS zip if not already cached."""
    if TEMP_ZIP.exists():
        print(f"  Using cached zip: {TEMP_ZIP}")
        return TEMP_ZIP
    print(f"  Downloading {CMS_ZIP_URL} ...")
    urlretrieve(CMS_ZIP_URL, TEMP_ZIP)
    print(f"  Saved to {TEMP_ZIP}")
    return TEMP_ZIP


def extract_xml(zip_path: Path) -> ET.Element:
    """Extract and parse the tabular XML from the zip."""
    print(f"  Extracting {XML_INNER_PATH} from zip ...")
    with zipfile.ZipFile(zip_path) as zf:
        # Try exact path first, then search
        names = zf.namelist()
        xml_name = None
        for name in names:
            if name == XML_INNER_PATH:
                xml_name = name
                break
        if xml_name is None:
            # Fallback: find any tabular xml
            for name in names:
                if re.search(r"tabular.*\.xml$", name, re.IGNORECASE):
                    xml_name = name
                    break
        if xml_name is None:
            raise FileNotFoundError(
                f"Could not find tabular XML in zip. Files: {names[:20]}"
            )
        print(f"  Found XML: {xml_name}")
        with zf.open(xml_name) as xf:
            tree = ET.parse(xf)
    return tree.getroot()


def parse_hierarchy(root: ET.Element) -> list[dict[str, str]]:
    """Parse the CMS XML into flat code rows with hierarchy info."""
    rows: list[dict[str, str]] = []

    for chapter in root.findall(".//chapter"):
        chapter_desc = _text(chapter.find("desc"))

        for section in chapter.findall(".//sectionRef"):
            # sectionRef doesn't contain codes; actual sections do
            pass

        for section in chapter.findall(".//section"):
            section_desc = _text(section.find("desc"))

            # Recursively collect all diag nodes under this section
            _collect_diags(section, chapter_desc, section_desc, "", rows)

    return rows


def _collect_diags(
    parent: ET.Element,
    chapter_desc: str,
    section_desc: str,
    parent_category_desc: str,
    rows: list[dict[str, str]],
) -> None:
    """Recursively collect diagnosis codes from nested <diag> elements."""
    for diag in parent.findall("diag"):
        code = _text(diag.find("name"))
        desc = _text(diag.find("desc"))

        # Determine category_desc: the 3-character level description
        raw_code = code.replace(".", "")
        if len(raw_code) == 3:
            category_desc = desc
        else:
            category_desc = parent_category_desc if parent_category_desc else desc

        # Add this code
        rows.append({
            "icd10_code": code,
            "description": desc.lower() if desc else "",
            "chapter_desc": chapter_desc,
            "section_desc": section_desc,
            "category_desc": category_desc.lower() if category_desc else "",
        })

        # Check for 7th character extensions
        seven_chr_def = diag.find("sevenChrDef")
        if seven_chr_def is not None:
            extensions = seven_chr_def.findall("extension")
            for ext in extensions:
                char = ext.get("char", "")
                ext_desc = ext.text or ""
                
                # For codes that need padding with X to reach 7 chars + extension
                padded_code = raw_code
                if len(padded_code) < 7:
                    padded_code = padded_code.ljust(7, "X")
                ext_code = padded_code + char
                
                # Format with dot
                if len(ext_code) > 3:
                    ext_code_dotted = ext_code[:3] + "." + ext_code[3:]
                else:
                    ext_code_dotted = ext_code

                full_desc = f"{desc.lower()}, {ext_desc.lower()}" if desc else ext_desc.lower()
                rows.append({
                    "icd10_code": ext_code_dotted,
                    "description": full_desc,
                    "chapter_desc": chapter_desc,
                    "section_desc": section_desc,
                    "category_desc": category_desc.lower() if category_desc else "",
                })

        # Recurse into child diags
        _collect_diags(diag, chapter_desc, section_desc, category_desc, rows)


def _text(el: ET.Element | None) -> str:
    """Safely get text from an XML element."""
    if el is None:
        return ""
    return (el.text or "").strip()


def deduplicate(rows: list[dict[str, str]]) -> list[dict[str, str]]:
    """Remove duplicate codes, keeping first occurrence."""
    seen: set[str] = set()
    unique: list[dict[str, str]] = []
    for row in rows:
        norm = row["icd10_code"].replace(".", "").upper()
        if norm not in seen:
            seen.add(norm)
            unique.append(row)
    return unique


def write_compressed_csv(rows: list[dict[str, str]]) -> None:
    """Write rows to a gzip-compressed CSV."""
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    print(f"  Writing {len(rows)} rows to {OUTPUT_FILE} ...")
    with gzip.open(OUTPUT_FILE, "wt", encoding="utf-8", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=COLUMNS)
        writer.writeheader()
        writer.writerows(rows)
    size_mb = OUTPUT_FILE.stat().st_size / (1024 * 1024)
    print(f"  Done! File size: {size_mb:.2f} MB")


def main() -> None:
    print("Building icdcodex bundled dataset ...")
    zip_path = download_zip()
    root = extract_xml(zip_path)
    rows = parse_hierarchy(root)
    print(f"  Parsed {len(rows)} raw rows")
    rows = deduplicate(rows)
    print(f"  After dedup: {len(rows)} unique codes")
    write_compressed_csv(rows)
    print("Done!")


if __name__ == "__main__":
    main()
