stack/dev/scripts/generate_ccw_docs.py

"""Parse the CCW FFS Claims Codebook PDF and generate structured docs.

Reads ``codebook-ffs-claims.pdf`` and produces one Python module per
variable containing the full documentation as presented in the PDF:
variable name, label, description, short/long names, type, length,
source, values, and comments.

Usage:
    python generate_ccw_docs.py [--pdf PATH] [--out DIR]

Defaults:
    --pdf  codebook-ffs-claims.pdf  (same directory)
    --out  ../src/ccw/docs/
"""

from __future__ import annotations

import argparse
import re
from dataclasses import dataclass
from pathlib import Path

import pdfplumber

# ── Footer pattern (appears on every page, must be stripped) ─────────────

FOOTER_RE = re.compile(r"^Chronic Conditions Warehouse\s+Virtual Research Data Center$")
FOOTER2_RE = re.compile(r"^Medicare FFS Claims.*Codebook.*V\s*\d+\.\d+\s+\d*$")
BACK_TO_TOC = "^ Back to TOC ^"

# ── Field label patterns in variable entries ─────────────────────────────

FIELD_LABELS = [
    "LABEL:",
    "DESCRIPTION:",
    "SHORT NAME:",
    "LONG NAME:",
    "TYPE:",
    "LENGTH:",
    "SOURCE:",
    "VALUES:",
    "COMMENT:",
]


@dataclass
class Variable:
    """A single parsed variable entry from the codebook."""

    name: str = ""
    label: str = ""
    description: str = ""
    short_name: str = ""
    long_name: str = ""
    type: str = ""
    length: str = ""
    source: str = ""
    values: str = ""
    comment: str = ""


# ── PDF text extraction ─────────────────────────────────────────────────


def extract_all_text(pdf_path: str) -> str:
    """Extract all text from the PDF, stripping page footers."""
    pdf = pdfplumber.open(pdf_path)
    all_lines: list[str] = []

    for page in pdf.pages:
        text = page.extract_text()
        if not text:
            continue
        for line in text.split("\n"):
            # Strip footer lines
            if FOOTER_RE.match(line.strip()):
                continue
            if FOOTER2_RE.match(line.strip()):
                continue
            all_lines.append(line)

    pdf.close()
    return "\n".join(all_lines)


# ── Variable entry parsing ──────────────────────────────────────────────


def find_variable_details_start(text: str) -> int:
    """Find the index where 'Variable Details' section begins."""
    marker = "Variable Details\n"
    idx = text.find(marker)
    if idx == -1:
        # Try alternate
        marker = "Variable Details"
        idx = text.find(marker)
    return idx


def split_into_entries(text: str) -> list[str]:
    """Split the variable details section into individual entries.

    Each entry ends with '^ Back to TOC ^'.
    """
    entries = text.split(BACK_TO_TOC)
    # Last chunk after final Back to TOC is empty/garbage
    return [e.strip() for e in entries if e.strip()]


def parse_entry(raw: str) -> Variable | None:
    """Parse a single variable entry into a Variable object."""
    lines = raw.split("\n")

    # Find the variable name — it's the line(s) before "LABEL:"
    label_idx = -1
    for i, line in enumerate(lines):
        if line.strip().startswith("LABEL:"):
            label_idx = i
            break

    if label_idx == -1:
        return None

    # Variable name is everything before LABEL:
    # Could be multiple lines for grouped variables like THRPY_CAP_IND_CD1-5
    name_lines = []
    for i in range(label_idx):
        line = lines[i].strip()
        if not line:
            continue
        # Skip the "Variable Details" header and intro text
        if line.startswith("Variable Details"):
            continue
        if line.startswith("This section of the codebook"):
            continue
        if line.startswith("Each entry contains"):
            continue
        name_lines.append(line)

    if not name_lines:
        return None

    var = Variable()
    var.name = "\n".join(name_lines)

    # Parse the remaining structured fields
    # Rejoin everything from LABEL: onward
    "\n".join(lines[label_idx:])

    # Extract each field by finding the label and collecting text
    # until the next label
    current_field = None
    current_text: list[str] = []

    for line in lines[label_idx:]:
        stripped = line.strip()

        # Check if this line starts a new field
        matched_field = None
        for fl in FIELD_LABELS:
            if stripped.startswith(fl):
                matched_field = fl
                break

        if matched_field:
            # Save previous field
            if current_field:
                _set_field(var, current_field, "\n".join(current_text).strip())

            current_field = matched_field
            # Get the rest of the line after the label
            rest = stripped[len(matched_field) :].strip()
            current_text = [rest] if rest else []
        else:
            # Continuation line
            if current_field:
                current_text.append(line.rstrip())

    # Save last field
    if current_field:
        _set_field(var, current_field, "\n".join(current_text).strip())

    return var


def _set_field(var: Variable, field_label: str, value: str) -> None:
    """Set a Variable field based on the field label string."""
    mapping = {
        "LABEL:": "label",
        "DESCRIPTION:": "description",
        "SHORT NAME:": "short_name",
        "LONG NAME:": "long_name",
        "TYPE:": "type",
        "LENGTH:": "length",
        "SOURCE:": "source",
        "VALUES:": "values",
        "COMMENT:": "comment",
    }
    attr = mapping.get(field_label)
    if attr:
        setattr(var, attr, value)


# ── Output generation ───────────────────────────────────────────────────


def variable_to_py(var: Variable) -> str:
    """Convert a Variable to a Python module with docstring content."""
    # Clean up the variable name for the module-level docstring
    primary_name = var.name.split("\n")[0].strip()

    lines = [
        '"""',
        f"Variable: {primary_name}",
        "",
    ]

    if var.label and var.label != "—":
        lines.append(f"Label: {var.label}")
        lines.append("")

    if var.description and var.description != "—":
        lines.append("Description")
        lines.append("-----------")
        lines.append(var.description)
        lines.append("")

    # Metadata block
    lines.append("Metadata")
    lines.append("--------")
    all_names = var.name.split("\n")
    if len(all_names) > 1:
        lines.append(f"Variables: {', '.join(n.strip() for n in all_names)}")
    if var.short_name:
        lines.append(f"Short SAS Name: {var.short_name}")
    if var.long_name:
        lines.append(f"Long SAS Name: {var.long_name}")
    if var.type:
        lines.append(f"Type: {var.type}")
    if var.length:
        lines.append(f"Length: {var.length}")
    if var.source:
        lines.append(f"Source: {var.source}")
    lines.append("")

    if var.values and var.values != "—":
        lines.append("Values")
        lines.append("------")
        lines.append(var.values)
        lines.append("")

    if var.comment and var.comment != "—":
        lines.append("Comment")
        lines.append("-------")
        lines.append(var.comment)
        lines.append("")

    lines.append('"""')
    lines.append("")

    # Add a simple dict for programmatic access
    lines.append(f"name = {primary_name!r}")
    lines.append(f"label = {var.label!r}")
    lines.append(f"short_name = {var.short_name!r}")
    lines.append(f"long_name = {var.long_name!r}")
    lines.append(f"type = {var.type!r}")
    lines.append(f"length = {var.length!r}")
    lines.append(f"source = {var.source!r}")
    lines.append("")

    return "\n".join(lines)


def sanitize_filename(name: str) -> str:
    """Convert a variable name to a safe Python module filename."""
    # Take first variable name if grouped
    primary = name.split("\n")[0].strip().lower()
    # Handle grouped names like CLM_POA_IND_SW1 through SW25
    return primary


# ── Main ────────────────────────────────────────────────────────────────


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate CCW codebook docs from PDF")
    parser.add_argument(
        "--pdf",
        default=str(
            Path(__file__).resolve().parents[2]
            / "dev"
            / "seeds"
            / "codebook-ffs-claims.pdf"
        ),
        help="Path to the CCW codebook PDF",
    )
    parser.add_argument(
        "--out",
        default=str(Path(__file__).resolve().parents[2] / "src" / "ccw" / "docs"),
        help="Output directory",
    )
    args = parser.parse_args()

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Extracting text from {args.pdf}...")
    full_text = extract_all_text(args.pdf)

    # Find where variable details start
    start_idx = find_variable_details_start(full_text)
    if start_idx == -1:
        print("ERROR: Could not find 'Variable Details' section")
        return

    details_text = full_text[start_idx:]

    print("Splitting into variable entries...")
    entries = split_into_entries(details_text)
    print(f"Found {len(entries)} raw entries")

    variables: list[Variable] = []
    for raw in entries:
        var = parse_entry(raw)
        if var and var.name:
            variables.append(var)

    print(f"Parsed {len(variables)} variables")

    # Write individual variable modules
    all_var_names: list[str] = []
    for var in variables:
        fname = sanitize_filename(var.name)
        module_path = out_dir / f"{fname}.py"
        module_path.write_text(variable_to_py(var))
        all_var_names.append(fname)

    # Write __init__.py with index
    init_lines = [
        '"""CCW FFS Claims Codebook — variable documentation.',
        "",
        "Auto-generated from codebook-ffs-claims.pdf.",
        f"Contains {len(variables)} variable entries.",
        '"""',
        "",
        "VARIABLES = [",
    ]
    for vname in sorted(all_var_names):
        init_lines.append(f'    "{vname}",')
    init_lines.append("]")
    init_lines.append("")
    (out_dir / "__init__.py").write_text("\n".join(init_lines))

    print(f"\nWrote {len(variables)} variable docs to {out_dir}/")
    print(f"Wrote {out_dir / '__init__.py'}")

    # Show summary
    types = {}
    for var in variables:
        t = var.type or "unknown"
        types[t] = types.get(t, 0) + 1
    print(f"\nType distribution: {types}")

    with_values = sum(1 for v in variables if v.values and v.values != "—")
    with_comment = sum(1 for v in variables if v.comment and v.comment != "—")
    print(f"Variables with values: {with_values}")
    print(f"Variables with comments: {with_comment}")


if __name__ == "__main__":
    main()