"""data loader for icdlookup.

handles lazy-loading of the built-in gzip-compressed icd-10 csv dataset,
as well as loading custom user-provided csv reference files.

the built-in dataset is sourced from official cms icd-10 files
(https://www.cms.gov/medicare/coding-billing/icd-10-codes).

design notes
------------
- only stdlib modules (csv, gzip, pathlib) are used here — no pandas required.
- the built-in dataset is parsed once on first access, then cached in-memory.
- custom reference csvs must contain at minimum ``icd10_code`` and
  ``description`` columns.
"""

from __future__ import annotations

import csv
import gzip
import io
from pathlib import Path
from typing import Any

_CACHE: dict[str, list[dict[str, str]]] = {}

_DATA_DIR = Path(__file__).resolve().parent / "_data"
_DEFAULT_DATASET = _DATA_DIR / "icd10_2026.csv.gz"


def normalise_code(code: str) -> str:
    """normalise an icd code for lookup: uppercase, strip whitespace & dots."""
    return code.strip().upper().replace(".", "")


def _add_dot(code: str) -> str:
    """return dotted form of a normalised code (e.g. 'A001' → 'A00.1')."""
    raw = code.replace(".", "")
    if len(raw) > 3:
        return raw[:3] + "." + raw[3:]
    return raw


def _load_gzip_csv(path: Path) -> list[dict[str, str]]:
    """read a gzip-compressed csv and return a list of row dicts."""
    with gzip.open(path, "rt", encoding="utf-8") as fh:
        reader = csv.DictReader(fh)
        return list(reader)


def _load_plain_csv(path: str | Path) -> list[dict[str, str]]:
    """read a plain csv and return a list of row dicts."""
    path = Path(path)
    if path.suffix == ".gz" or path.name.endswith(".csv.gz"):
        return _load_gzip_csv(path)
    with open(path, encoding="utf-8", newline="") as fh:
        reader = csv.DictReader(fh)
        return list(reader)


def load_data(reference: str | Path | None = None) -> list[dict[str, str]]:
    """return the reference dataset as a list of row dicts.

    parameters
    ----------
    reference : str | Path | None
        path to a custom csv (plain or ``.csv.gz``).  must contain at least
        ``icd10_code`` and ``description`` columns.  when *None* the built-in
        2026 dataset is used.

    returns
    -------
    list[dict[str, str]]
        each dict represents one row from the csv.
    """
    cache_key = str(reference) if reference else "__builtin__"

    if cache_key in _CACHE:
        return _CACHE[cache_key]

    if reference is None:
        if not _DEFAULT_DATASET.exists():
            raise FileNotFoundError(
                f"built-in dataset not found at {_DEFAULT_DATASET}. "
                "please reinstall the package or provide a custom reference csv."
            )
        rows = _load_gzip_csv(_DEFAULT_DATASET)
    else:
        rows = _load_plain_csv(reference)

    # validate required columns
    if rows:
        first = rows[0]
        for col in ("icd10_code", "description"):
            if col not in first:
                raise ValueError(
                    f"reference csv is missing required column '{col}'. "
                    f"found columns: {list(first.keys())}"
                )

    _CACHE[cache_key] = rows
    return rows


def build_index(
    reference: str | Path | None = None,
) -> dict[str, dict[str, str]]:
    """build a normalised-code → row-dict index for fast O(1) lookups.

    the index is also cached so repeat calls are free.
    """
    idx_key = ("__idx__", str(reference) if reference else "__builtin__")
    # we stash it in the global cache under a tuple key (won't collide with
    # plain string keys used by load_data).
    idx_key_str = repr(idx_key)
    if idx_key_str in _CACHE:
        return _CACHE[idx_key_str]  # type: ignore[return-value]

    rows = load_data(reference)
    index: dict[str, dict[str, str]] = {}
    for row in rows:
        norm = normalise_code(row["icd10_code"])
        # first occurrence wins (shouldn't have dupes, but be safe)
        if norm not in index:
            index[norm] = row

    _CACHE[idx_key_str] = index  # type: ignore[assignment]
    return index


def clear_cache() -> None:
    """clear all cached data (useful for testing or memory management)."""
    _CACHE.clear()
