"""Validate data/climate-data.csv before the browser app loads it.

Checks file format, columns, county keys, the county geometry join, per-metric
types and ranges, allowed blank values, and cross-field consistency. Exits with
status 1 when any check fails.

Run:
    .venv\\Scripts\\python.exe scripts\\check_climate_data.py
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, FrozenSet, List, Optional, Sequence, Tuple

PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CLIMATE_CSV = PROJECT_ROOT / "data" / "climate-data.csv"
DEFAULT_COUNTIES_GEOJSON = PROJECT_ROOT / "data" / "geojson-counties-fips.json"
DEFAULT_METRIC_SOURCES = PROJECT_ROOT / "data" / "metric_sources.json"

FIPS_PATTERN = re.compile(r"\d{5}")
STATE_PATTERN = re.compile(r"[A-Z]{2}")

# The distinct-value check only runs on full-size files, not small fixtures.
COLLAPSE_CHECK_MIN_ROWS = 100

KOPPEN_CODES = (
    "Af", "Am", "Aw", "BWh", "BWk", "BSh", "BSk",
    "Csa", "Csb", "Csc", "Cwa", "Cwb", "Cwc", "Cfa", "Cfb", "Cfc",
    "Dsa", "Dsb", "Dsc", "Dsd", "Dwa", "Dwb", "Dwc", "Dwd",
    "Dfa", "Dfb", "Dfc", "Dfd", "ET", "EF",
)
# Counties with no predominant class; see docs/filter-calculations.md, section 1.
MIXED_KOPPEN_CLASS = "Mixed"

MONTH_NAMES = (
    "January", "February", "March", "April", "May", "June",
    "July", "August", "September", "October", "November", "December",
)

# NOAA nClimGrid and gridMET cover the contiguous U.S. only.
OUTSIDE_CONUS = frozenset({"AK", "HI", "PR"})
# NSRDB summaries were not requested for Puerto Rico.
NSRDB_NOT_REQUESTED = frozenset({"PR"})
# NOAA county daily files do not include Lexington city, VA.
NOAA_DAILY_MISSING_FIPS = frozenset({"51678"})


@dataclass(frozen=True)
class MetricRule:
    """Validation rule for one metric column."""

    kind: str
    minimum: Optional[float] = None
    maximum: Optional[float] = None
    integer: bool = False
    categories: Tuple[str, ...] = ()
    blank_states: FrozenSet[str] = frozenset()
    blank_fips: FrozenSet[str] = frozenset()
    min_distinct: int = 2


def _numeric(
    minimum: float,
    maximum: float,
    *,
    integer: bool = False,
    blank_states: FrozenSet[str] = frozenset(),
    blank_fips: FrozenSet[str] = frozenset(),
) -> MetricRule:
    """Build a numeric rule with a plausible physical range."""
    return MetricRule(
        "numeric",
        minimum=minimum,
        maximum=maximum,
        integer=integer,
        blank_states=blank_states,
        blank_fips=blank_fips,
        min_distinct=10,
    )


def _categorical(categories: Sequence[str], *, blank_states: FrozenSet[str] = frozenset()) -> MetricRule:
    """Build a categorical rule limited to an allowed value list."""
    return MetricRule("categorical", categories=tuple(categories), blank_states=blank_states)


# Ranges are physical plausibility limits, deliberately wider than the current
# data. They are not the app's color-scale bounds.
METRIC_RULES: Dict[str, MetricRule] = {
    "koppenZone": _categorical(KOPPEN_CODES + (MIXED_KOPPEN_CLASS,)),
    "avgTempF": _numeric(20, 85, blank_states=OUTSIDE_CONUS),
    "avgDiurnalTempRangeF": _numeric(5, 40, blank_states=OUTSIDE_CONUS, blank_fips=NOAA_DAILY_MISSING_FIPS),
    "annualPrecipIn": _numeric(1, 150, blank_states=OUTSIDE_CONUS),
    "seasonalityIndex": _numeric(0, 100, integer=True, blank_states=OUTSIDE_CONUS),
    "wettestPrecipMonth": _categorical(MONTH_NAMES, blank_states=OUTSIDE_CONUS),
    "driestPrecipMonth": _categorical(MONTH_NAMES, blank_states=OUTSIDE_CONUS),
    "absoluteExtremeDays": _numeric(0, 366, blank_states=OUTSIDE_CONUS, blank_fips=NOAA_DAILY_MISSING_FIPS),
    "meanDailyGlobalHorizontalRadiationKwhM2Day": _numeric(1.5, 7.5, blank_states=NSRDB_NOT_REQUESTED),
    "clearSkyGhiReductionIndex": _numeric(0, 1, blank_states=NSRDB_NOT_REQUESTED),
    "avgSummerSpecificHumidityGKg": _numeric(1, 25, blank_states=OUTSIDE_CONUS),
    "humidHeatDays": _numeric(0, 366, blank_states=OUTSIDE_CONUS),
}

IDENTITY_COLUMNS = ("countyFips", "countyName", "state")
# Stripe classes the app draws for Mixed Köppen counties; blank otherwise.
KOPPEN_STRIPE_COLUMNS = ("koppenPrimaryClass", "koppenSecondaryClass")
AUDIT_COLUMNS = ("humidHeatSourceFips", "humidHeatFipsAdjustment", "source")
EXPECTED_COLUMNS = IDENTITY_COLUMNS + tuple(METRIC_RULES) + KOPPEN_STRIPE_COLUMNS + AUDIT_COLUMNS

CHECK_FORMAT = "file format"
CHECK_COLUMNS = "columns"
CHECK_COUNTIES = "county keys"
CHECK_GEOMETRY = "map geometry join"
CHECK_BLANKS = "blank values"
CHECK_VALUES = "value types and ranges"
CHECK_CROSS = "cross-field consistency"
CHECK_SOURCES = "metric sources file"


@dataclass
class Report:
    """Problems found per check, in the order checks ran."""

    checks: Dict[str, List[str]] = field(default_factory=dict)
    notes: List[str] = field(default_factory=list)
    row_count: int = 0
    column_count: int = 0

    def start(self, check: str) -> None:
        self.checks.setdefault(check, [])

    def add(self, check: str, message: str) -> None:
        self.checks.setdefault(check, []).append(message)

    @property
    def problem_count(self) -> int:
        return sum(len(messages) for messages in self.checks.values())

    @property
    def ok(self) -> bool:
        return self.problem_count == 0


def read_climate_csv(path: Path) -> Tuple[List[str], List[List[str]]]:
    """Read the raw header and rows without DictReader's silent padding."""
    with path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.reader(handle)
        headers = next(reader, [])
        rows = [row for row in reader if row]
    return headers, rows


def check_structure(
    report: Report, headers: List[str], raw_rows: List[List[str]]
) -> Tuple[List[str], List[Dict[str, str]]]:
    """Check encoding, header names, and row widths; return rows as dicts."""
    report.start(CHECK_FORMAT)
    report.start(CHECK_COLUMNS)

    if headers and headers[0].startswith("\ufeff"):
        report.add(CHECK_FORMAT, "file starts with a UTF-8 BOM; the app would not find the first column")
        headers = [headers[0].lstrip("\ufeff")] + headers[1:]

    seen = set()
    for name in headers:
        if name in seen:
            report.add(CHECK_COLUMNS, f"duplicate column {name!r}")
        seen.add(name)

    for name in EXPECTED_COLUMNS:
        if name not in seen:
            report.add(CHECK_COLUMNS, f"missing column {name!r}")
    for name in headers:
        if name not in EXPECTED_COLUMNS:
            report.add(CHECK_COLUMNS, f"unexpected column {name!r} (add it to check_climate_data.py if intended)")

    rows: List[Dict[str, str]] = []
    for index, raw in enumerate(raw_rows):
        if len(raw) != len(headers):
            report.add(CHECK_FORMAT, f"line {index + 2}: {len(raw)} fields, expected {len(headers)}")
            continue
        rows.append(dict(zip(headers, raw)))

    report.row_count = len(raw_rows)
    report.column_count = len(headers)
    return headers, rows


def check_counties(report: Report, rows: List[Dict[str, str]]) -> None:
    """Check FIPS format and uniqueness, names, and state/FIPS agreement."""
    report.start(CHECK_COUNTIES)
    seen = set()
    prefix_states: Dict[str, set] = {}
    state_prefixes: Dict[str, set] = {}

    for row in rows:
        fips = row.get("countyFips") or ""
        if not FIPS_PATTERN.fullmatch(fips):
            report.add(CHECK_COUNTIES, f"invalid countyFips {fips!r}")
            continue
        if fips in seen:
            report.add(CHECK_COUNTIES, f"duplicate countyFips {fips}")
        seen.add(fips)

        if not (row.get("countyName") or "").strip():
            report.add(CHECK_COUNTIES, f"{fips}: countyName is blank")

        state = row.get("state") or ""
        if not STATE_PATTERN.fullmatch(state):
            report.add(CHECK_COUNTIES, f"{fips}: invalid state {state!r}")
            continue
        prefix_states.setdefault(fips[:2], set()).add(state)
        state_prefixes.setdefault(state, set()).add(fips[:2])

    for prefix, states in sorted(prefix_states.items()):
        if len(states) > 1:
            report.add(CHECK_COUNTIES, f"state FIPS {prefix} is labeled as {sorted(states)}")
    for state, prefixes in sorted(state_prefixes.items()):
        if len(prefixes) > 1:
            report.add(CHECK_COUNTIES, f"state {state} spans FIPS prefixes {sorted(prefixes)}")


def _feature_fips(feature: dict) -> str:
    """Return the 5-digit county FIPS stored on a GeoJSON feature."""
    props = feature.get("properties") or {}
    raw = feature.get("id") or props.get("id") or str(props.get("GEO_ID") or "")[-5:]
    text = str(raw or "").strip()
    return text.zfill(5) if text.isdigit() else text


def check_geometry(report: Report, rows: List[Dict[str, str]], geojson_path: Path) -> None:
    """Check that CSV counties and map polygons match one-to-one."""
    report.start(CHECK_GEOMETRY)
    if not geojson_path.exists():
        report.add(CHECK_GEOMETRY, f"county geometry file not found: {geojson_path}")
        return

    features = json.loads(geojson_path.read_text(encoding="utf-8")).get("features", [])
    geo_fips = set()
    for feature in features:
        fips = _feature_fips(feature)
        if not FIPS_PATTERN.fullmatch(fips):
            report.add(CHECK_GEOMETRY, f"map feature without a valid county FIPS: {fips!r}")
            continue
        if fips in geo_fips:
            report.add(CHECK_GEOMETRY, f"map has duplicate polygons for {fips}")
        geo_fips.add(fips)
        state_fips = str((feature.get("properties") or {}).get("STATE") or "").strip()
        if state_fips and state_fips.zfill(2) != fips[:2]:
            report.add(CHECK_GEOMETRY, f"map feature {fips} has STATE {state_fips}")

    csv_fips = {row.get("countyFips") or "" for row in rows}
    for fips in sorted(csv_fips - geo_fips):
        report.add(CHECK_GEOMETRY, f"{fips} is in the CSV but has no map polygon")
    for fips in sorted(geo_fips - csv_fips):
        report.add(CHECK_GEOMETRY, f"{fips} has a map polygon but no CSV row")


def check_metrics(report: Report, headers: List[str], rows: List[Dict[str, str]]) -> None:
    """Check each metric's blanks, types, ranges, and categories."""
    report.start(CHECK_BLANKS)
    report.start(CHECK_VALUES)

    for key, rule in METRIC_RULES.items():
        if key not in headers:
            continue
        present: List[object] = []
        for row in rows:
            fips = row.get("countyFips") or "?"
            state = row.get("state") or ""
            value = row.get(key) or ""

            if value == "":
                if state not in rule.blank_states and fips not in rule.blank_fips:
                    report.add(CHECK_BLANKS, f"{fips} ({state}): {key} is blank")
                continue
            if value != value.strip():
                report.add(CHECK_VALUES, f"{fips}: {key}={value!r} has surrounding whitespace")
                continue

            if rule.kind == "categorical":
                if value in rule.categories:
                    present.append(value)
                else:
                    report.add(CHECK_VALUES, f"{fips}: {key}={value!r} is not an allowed category")
                continue

            try:
                number = float(value)
            except ValueError:
                report.add(CHECK_VALUES, f"{fips}: {key}={value!r} is not a number")
                continue
            if not math.isfinite(number):
                report.add(CHECK_VALUES, f"{fips}: {key}={value!r} is not finite")
                continue
            if rule.integer and not number.is_integer():
                report.add(CHECK_VALUES, f"{fips}: {key}={value} should be a whole number")
            if not rule.minimum <= number <= rule.maximum:
                report.add(CHECK_VALUES, f"{fips}: {key}={value} outside {rule.minimum:g}..{rule.maximum:g}")
            present.append(number)

        if len(rows) >= COLLAPSE_CHECK_MIN_ROWS and len(set(present)) < rule.min_distinct:
            report.add(
                CHECK_VALUES,
                f"{key} has only {len(set(present))} distinct values; the column may have been overwritten",
            )


def check_cross_fields(report: Report, rows: List[Dict[str, str]]) -> None:
    """Check relationships between columns within each row."""
    report.start(CHECK_CROSS)
    for row in rows:
        fips = row.get("countyFips") or "?"

        wet = row.get("wettestPrecipMonth") or ""
        dry = row.get("driestPrecipMonth") or ""
        if bool(wet) != bool(dry):
            report.add(CHECK_CROSS, f"{fips}: only one of wettestPrecipMonth/driestPrecipMonth is set")
        elif wet and wet == dry:
            report.add(CHECK_CROSS, f"{fips}: wettest and driest month are both {wet}")

        heat = row.get("humidHeatDays") or ""
        heat_source = row.get("humidHeatSourceFips") or ""
        if bool(heat) != bool(heat_source):
            report.add(CHECK_CROSS, f"{fips}: humidHeatDays and humidHeatSourceFips must both be set or both blank")
        if heat_source and not FIPS_PATTERN.fullmatch(heat_source):
            report.add(CHECK_CROSS, f"{fips}: invalid humidHeatSourceFips {heat_source!r}")
        elif heat_source and heat_source != fips and not (row.get("humidHeatFipsAdjustment") or "").strip():
            report.add(CHECK_CROSS, f"{fips}: uses proxy county {heat_source} without a humidHeatFipsAdjustment note")

        zone = row.get("koppenZone") or ""
        primary = row.get("koppenPrimaryClass") or ""
        secondary = row.get("koppenSecondaryClass") or ""
        if zone == MIXED_KOPPEN_CLASS:
            if not (primary and secondary):
                report.add(CHECK_CROSS, f"{fips}: Mixed Köppen county needs koppenPrimaryClass and koppenSecondaryClass")
            elif primary not in KOPPEN_CODES or secondary not in KOPPEN_CODES:
                report.add(CHECK_CROSS, f"{fips}: invalid Köppen stripe classes {primary!r}/{secondary!r}")
            elif primary == secondary:
                report.add(CHECK_CROSS, f"{fips}: Köppen stripe classes are both {primary}")
        elif primary or secondary:
            report.add(CHECK_CROSS, f"{fips}: Köppen stripe classes are set but koppenZone is {zone or 'blank'}, not Mixed")

        if not (row.get("source") or "").strip():
            report.add(CHECK_CROSS, f"{fips}: source is blank")


def check_metric_sources(report: Report, path: Path) -> None:
    """Check that the metric metadata file parses and names real metrics."""
    report.start(CHECK_SOURCES)
    if not path.exists():
        report.add(CHECK_SOURCES, f"{path.name} not found")
        return
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        report.add(CHECK_SOURCES, f"{path.name} is not valid JSON: {exc}")
        return
    if not isinstance(data, dict) or not isinstance(data.get("metrics"), dict):
        report.add(CHECK_SOURCES, f'{path.name} must be a JSON object with a "metrics" object')
        return

    documented = set(data["metrics"])
    for key in sorted(documented - set(METRIC_RULES)):
        report.add(CHECK_SOURCES, f"{path.name} describes unknown metric {key!r}")
    covered = len(documented & set(METRIC_RULES))
    report.notes.append(f"{path.name} documents {covered} of {len(METRIC_RULES)} metrics.")


def run_checks(csv_path: Path, geojson_path: Path, metric_sources_path: Path) -> Report:
    """Run every check and return the combined report."""
    report = Report()
    headers, raw_rows = read_climate_csv(csv_path)
    headers, rows = check_structure(report, headers, raw_rows)
    check_counties(report, rows)
    check_geometry(report, rows, geojson_path)
    check_metrics(report, headers, rows)
    check_cross_fields(report, rows)
    check_metric_sources(report, metric_sources_path)
    return report


def print_report(report: Report, csv_path: Path, max_examples: int) -> None:
    """Print a pass/fail line per check with example problems."""
    print(f"Checked {csv_path}: {report.row_count} rows, {report.column_count} columns")
    for check, messages in report.checks.items():
        status = "PASS" if not messages else f"FAIL ({len(messages)})"
        print(f"  {status:<11}{check}")
        for message in messages[:max_examples]:
            print(f"      - {message}")
        if len(messages) > max_examples:
            print(f"      ... and {len(messages) - max_examples} more")
    for note in report.notes:
        print(f"  note: {note}")
    print("Result: PASS" if report.ok else f"Result: FAIL ({report.problem_count} problems)")


def parse_args() -> argparse.Namespace:
    """Define and parse command-line options for this checker."""
    parser = argparse.ArgumentParser(description="Validate the browser app's county climate CSV.")
    parser.add_argument("--csv", type=Path, default=DEFAULT_CLIMATE_CSV, help="Climate data CSV path.")
    parser.add_argument("--geojson", type=Path, default=DEFAULT_COUNTIES_GEOJSON, help="County GeoJSON path.")
    parser.add_argument("--metric-sources", type=Path, default=DEFAULT_METRIC_SOURCES, help="Metric metadata JSON path.")
    parser.add_argument("--max-examples", type=int, default=10, help="Problems to print per failing check.")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if not args.csv.exists():
        print(f"Climate data CSV not found: {args.csv}", file=sys.stderr)
        return 2
    report = run_checks(args.csv, args.geojson, args.metric_sources)
    print_report(report, args.csv, args.max_examples)
    return 0 if report.ok else 1


if __name__ == "__main__":
    sys.exit(main())