"""
Mortality/Life table utilities.

Helpers to query `life_tables` and `number_tables` by age/month and
return values filtered by sex/race using compact codes:
  - sex: M, F, A (All)
  - race: W (White), B (Black), H (Hispanic), A (All)

Column naming in tables follows the pattern:
  - LifeTable: le_{race}{sex}, na_{race}{sex}
  - NumberTable: na_{race}{sex}

Examples:
  - race=W, sex=M => suffix "wm" (columns `le_wm`, `na_wm`)
  - race=A, sex=F => suffix "af" (columns `le_af`, `na_af`)
  - race=H, sex=A => suffix "ha" (columns `le_ha`, `na_ha`)
"""

from __future__ import annotations

from typing import Dict, Optional, Tuple
from sqlalchemy.orm import Session

from app.models.pensions import LifeTable, NumberTable


_RACE_MAP: Dict[str, str] = {
    "W": "w",  # White
    "B": "b",  # Black
    "H": "h",  # Hispanic
    "A": "a",  # All races
}

_SEX_MAP: Dict[str, str] = {
    "M": "m",
    "F": "f",
    "A": "a",  # All sexes
}


class InvalidCodeError(ValueError):
    pass


def _normalize_codes(sex: str, race: str) -> Tuple[str, str, str]:
    """Validate/normalize sex and race to construct the column suffix.

    Returns (suffix, sex_u, race_u) where suffix is lowercase like "wm".
    Raises InvalidCodeError on invalid inputs.
    """
    sex_u = (sex or "").strip().upper()
    race_u = (race or "").strip().upper()
    if sex_u not in _SEX_MAP:
        raise InvalidCodeError(f"Invalid sex code '{sex}'. Expected one of: {', '.join(_SEX_MAP.keys())}")
    if race_u not in _RACE_MAP:
        raise InvalidCodeError(f"Invalid race code '{race}'. Expected one of: {', '.join(_RACE_MAP.keys())}")
    return _RACE_MAP[race_u] + _SEX_MAP[sex_u], sex_u, race_u


def get_life_values(
    db: Session,
    *,
    age: int,
    sex: str,
    race: str,
) -> Optional[Dict[str, Optional[float]]]:
    """Return life table LE and NA values for a given age, sex, and race.

    Returns dict: {"age": int, "sex": str, "race": str, "le": float|None, "na": float|None}
    Returns None if the age row does not exist.
    Raises InvalidCodeError for invalid codes.
    """
    suffix, sex_u, race_u = _normalize_codes(sex, race)
    row: Optional[LifeTable] = db.query(LifeTable).filter(LifeTable.age == age).first()
    if not row:
        return None

    le_col = f"le_{suffix}"
    na_col = f"na_{suffix}"
    le_val = getattr(row, le_col, None)
    na_val = getattr(row, na_col, None)

    return {
        "age": int(age),
        "sex": sex_u,
        "race": race_u,
        "le": float(le_val) if le_val is not None else None,
        "na": float(na_val) if na_val is not None else None,
    }


def get_number_value(
    db: Session,
    *,
    month: int,
    sex: str,
    race: str,
) -> Optional[Dict[str, Optional[float]]]:
    """Return number table NA value for a given month, sex, and race.

    Returns dict: {"month": int, "sex": str, "race": str, "na": float|None}
    Returns None if the month row does not exist.
    Raises InvalidCodeError for invalid codes.
    """
    suffix, sex_u, race_u = _normalize_codes(sex, race)
    row: Optional[NumberTable] = db.query(NumberTable).filter(NumberTable.month == month).first()
    if not row:
        return None

    na_col = f"na_{suffix}"
    na_val = getattr(row, na_col, None)

    return {
        "month": int(month),
        "sex": sex_u,
        "race": race_u,
        "na": float(na_val) if na_val is not None else None,
    }


__all__ = [
    "InvalidCodeError",
    "get_life_values",
    "get_number_value",
]