"""Inventory NetCDF inputs and report metadata, integrity and compatibility.

The default configuration scans values only for the small observation and
mask files.  ECMWF raw files are inspected metadata-only unless explicitly
configured otherwise, avoiding an accidental full scan of a large archive.
"""
from __future__ import annotations

import argparse
import json
import logging
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Iterable

import numpy as np
import xarray as xr

from .config import PROJECT_ROOT, project_path


LOG = logging.getLogger(__name__)
COORDINATE_ALIASES = {
    "longitude": ("longitude", "lon"),
    "latitude": ("latitude", "lat"),
    "time": ("valid_time", "time"),
    "forecast_reference_time": ("forecast_reference_time", "init_time"),
    "lead": ("forecast_period", "step", "lead"),
    "member": ("number", "member", "realization"),
}


def _json_value(value: Any) -> Any:
    """Convert NumPy, datetime and timedelta scalars to JSON-safe values."""
    if isinstance(value, (datetime, np.datetime64, np.timedelta64)):
        return str(value)
    if isinstance(value, np.generic):
        value = value.item()
    if isinstance(value, bytes):
        return value.decode("utf-8", errors="replace")
    if isinstance(value, (str, int, float, bool)) or value is None:
        return value
    return str(value)


def _attrs(attrs: dict[str, Any]) -> dict[str, Any]:
    return {str(key): _json_value(value) for key, value in attrs.items()}


def _coordinate_name(ds: xr.Dataset, canonical: str) -> str | None:
    for name in COORDINATE_ALIASES[canonical]:
        if name in ds.coords or name in ds.variables:
            return name
    return None


def _numeric_differences(values: np.ndarray) -> np.ndarray | None:
    if values.size < 2:
        return np.array([], dtype=float)
    if np.issubdtype(values.dtype, np.datetime64):
        return np.diff(values.astype("datetime64[ns]")).astype("timedelta64[ns]").astype(np.int64)
    if np.issubdtype(values.dtype, np.timedelta64):
        return np.diff(values.astype("timedelta64[ns]")).astype(np.int64)
    if np.issubdtype(values.dtype, np.number):
        return np.diff(values.astype(float))
    return None


def _coordinate_summary(da: xr.DataArray) -> dict[str, Any]:
    values = np.asarray(da.values).reshape(-1)
    result: dict[str, Any] = {
        "name": da.name,
        "dims": list(da.dims),
        "size": int(values.size),
        "dtype": str(da.dtype),
        "attrs": _attrs(da.attrs),
        "encoding": _attrs({key: value for key, value in da.encoding.items() if key in {"units", "calendar", "dtype", "_FillValue"}}),
    }
    if not values.size:
        return result
    result["first"] = _json_value(values[0])
    result["last"] = _json_value(values[-1])
    differences = _numeric_differences(values)
    if differences is None:
        result["duplicate_count"] = int(values.size - np.unique(values).size)
        return result
    finite = differences[np.isfinite(differences)]
    result["duplicate_count"] = int(np.count_nonzero(finite == 0))
    result["monotonic"] = "increasing" if np.all(finite > 0) else "decreasing" if np.all(finite < 0) else "not_monotonic"
    if finite.size:
        absolute = np.abs(finite)
        unique, counts = np.unique(absolute, return_counts=True)
        modal = unique[np.argmax(counts)]
        result["modal_step"] = _json_value(
            np.timedelta64(int(modal), "ns")
            if np.issubdtype(values.dtype, (np.datetime64, np.timedelta64))
            else float(modal)
        )
        result["irregular_step_count"] = int(np.count_nonzero(absolute != modal))
        ratios = absolute / modal if modal else np.ones_like(absolute)
        result["estimated_missing_steps"] = int(np.sum(np.maximum(np.rint(ratios).astype(int) - 1, 0)))
        result["minimum"] = _json_value(np.nanmin(values))
        result["maximum"] = _json_value(np.nanmax(values))
    return result


def _variable_summary(da: xr.DataArray, statistics: bool) -> dict[str, Any]:
    result: dict[str, Any] = {
        "dims": list(da.dims),
        "shape": [int(size) for size in da.shape],
        "dtype": str(da.dtype),
        "attrs": _attrs(da.attrs),
        "encoding": _attrs({key: value for key, value in da.encoding.items() if key in {"units", "calendar", "dtype", "_FillValue", "missing_value"}}),
    }
    if statistics:
        values = np.asarray(da.values)
        if np.issubdtype(values.dtype, np.number):
            finite = np.isfinite(values)
            result["value_statistics"] = {
                "total_count": int(values.size),
                "finite_count": int(finite.sum()),
                "missing_count": int(values.size - finite.sum()),
                "minimum": float(np.nanmin(values)) if finite.any() else None,
                "maximum": float(np.nanmax(values)) if finite.any() else None,
                "mean": float(np.nanmean(values)) if finite.any() else None,
            }
    return result


def inventory_file(path: Path, expected_variables: Iterable[str] = (), statistics: bool = False) -> dict[str, Any]:
    """Return a JSON-serializable inventory for one NetCDF file."""
    report: dict[str, Any] = {
        "path": str(path),
        "size_bytes": path.stat().st_size,
        "statistics_scanned": statistics,
        "status": "ok",
    }
    try:
        with xr.open_dataset(path, decode_times=True) as ds:
            actual = list(ds.data_vars)
            expected = list(expected_variables)
            report.update({
                "dimensions": {name: int(size) for name, size in ds.sizes.items()},
                "data_variables": actual,
                "missing_expected_variables": sorted(set(expected) - set(actual)),
                "unexpected_variables": sorted(set(actual) - set(expected)) if expected else [],
                "global_attributes": _attrs(ds.attrs),
                "coordinates": {},
                "variables": {},
            })
            for canonical in COORDINATE_ALIASES:
                name = _coordinate_name(ds, canonical)
                if name is not None:
                    report["coordinates"][canonical] = _coordinate_summary(ds[name])
            for name in actual:
                report["variables"][name] = _variable_summary(ds[name], statistics)
    except Exception as exc:  # Report file-level failures without hiding them.
        report["status"] = "error"
        report["error"] = f"{type(exc).__name__}: {exc}"
    return report


def _grid_signature(report: dict[str, Any]) -> tuple[Any, ...] | None:
    coords = report.get("coordinates", {})
    lon, lat = coords.get("longitude"), coords.get("latitude")
    if not lon or not lat:
        return None
    return (lon.get("size"), lon.get("first"), lon.get("last"), lon.get("modal_step"),
            lat.get("size"), lat.get("first"), lat.get("last"), lat.get("modal_step"))


def _temporal_extent(report: dict[str, Any]) -> tuple[str, str] | None:
    coords = report.get("coordinates", {})
    coordinate = coords.get("time") or coords.get("forecast_reference_time")
    if coordinate and coordinate.get("first") is not None and coordinate.get("last") is not None:
        return str(coordinate["first"]), str(coordinate["last"])
    return None


def _initial_month(report: dict[str, Any]) -> str | None:
    """Return the YYYY-MM forecast initialization represented by one file."""
    coordinate = report.get("coordinates", {}).get("forecast_reference_time")
    if not coordinate or coordinate.get("size") != 1 or coordinate.get("first") is None:
        return None
    try:
        value = np.datetime64(str(coordinate["first"]), "M")
    except (TypeError, ValueError):
        return None
    return str(value)


def _initialization_summary(files: list[dict[str, Any]]) -> dict[str, Any] | None:
    months = [month for report in files if report.get("status") == "ok"
              if (month := _initial_month(report)) is not None]
    if not months:
        return None
    values = np.asarray(months, dtype="datetime64[M]")
    unique = np.unique(values)
    expected = np.arange(unique[0], unique[-1] + np.timedelta64(1, "M"), dtype="datetime64[M]")
    missing = np.setdiff1d(expected, unique)
    return {
        "count": int(values.size),
        "unique_count": int(unique.size),
        "start": str(unique[0]),
        "end": str(unique[-1]),
        "duplicate_months": [str(value) for value in unique if np.count_nonzero(values == value) > 1],
        "missing_months": [str(value) for value in missing],
        "continuous": not missing.size and unique.size == values.size,
    }


def summarize_collection(files: list[dict[str, Any]]) -> dict[str, Any]:
    """Summarize consistency and temporal coverage for a file collection."""
    ok = [report for report in files if report.get("status") == "ok"]
    grids = [_grid_signature(report) for report in ok]
    grids = [grid for grid in grids if grid is not None]
    variable_sets = [tuple(sorted(report.get("data_variables", []))) for report in ok]
    extents = [extent for report in ok if (extent := _temporal_extent(report))]
    missing_variables = Counter(
        variable for report in ok for variable in report.get("missing_expected_variables", [])
    )
    result = {
        "file_count": len(files),
        "readable_file_count": len(ok),
        "error_file_count": len(files) - len(ok),
        "total_size_bytes": sum(int(report.get("size_bytes", 0)) for report in files),
        "grid_consistent": len(set(grids)) <= 1 if grids else None,
        "variable_set_consistent": len(set(variable_sets)) <= 1 if variable_sets else None,
        "missing_expected_variable_file_counts": dict(sorted(missing_variables.items())),
        "time_start": min((extent[0] for extent in extents), default=None),
        "time_end": max((extent[1] for extent in extents), default=None),
    }
    initialization = _initialization_summary(files)
    if initialization is not None:
        result["initialization_months"] = initialization
    return result


def _paired_initialization_report(
    datasets: list[dict[str, Any]], specifications: Iterable[dict[str, Any]],
) -> list[dict[str, Any]]:
    by_name = {dataset["name"]: dataset for dataset in datasets}
    reports = []
    for specification in specifications:
        names = list(specification["datasets"])
        month_sets = {}
        for name in names:
            dataset = by_name.get(name)
            if dataset is None:
                month_sets[name] = set()
                continue
            month_sets[name] = {
                month for file_report in dataset["files"]
                if (month := _initial_month(file_report)) is not None
            }
        union = set().union(*month_sets.values()) if month_sets else set()
        reports.append({
            "name": specification.get("name", "__".join(names)),
            "datasets": names,
            "match": bool(month_sets) and all(months == union for months in month_sets.values()),
            "missing_by_dataset": {
                name: sorted(union - months) for name, months in month_sets.items() if union - months
            },
        })
    return reports


def _expand_paths(patterns: Iterable[str]) -> list[Path]:
    paths: list[Path] = []
    for pattern in patterns:
        candidate = Path(pattern)
        if candidate.is_absolute():
            parent, name = candidate.parent, candidate.name
            matches = sorted(parent.glob(name))
        else:
            matches = sorted(PROJECT_ROOT.glob(pattern))
        paths.extend(path for path in matches if path.is_file())
    return list(dict.fromkeys(paths))


def build_inventory(config_path: str | Path = "config/inventory.json", selected: set[str] | None = None) -> dict[str, Any]:
    source = project_path(config_path)
    with source.open(encoding="utf-8") as handle:
        config = json.load(handle)
    datasets = []
    for specification in config["datasets"]:
        name = specification["name"]
        if selected and name not in selected:
            continue
        paths = _expand_paths(specification["paths"])
        LOG.info("Inventory %s: %d file(s)", name, len(paths))
        files = [inventory_file(path, specification.get("expected_variables", ()), specification.get("statistics", False)) for path in paths]
        datasets.append({
            "name": name,
            "patterns": specification["paths"],
            "expected_variables": specification.get("expected_variables", []),
            "statistics": specification.get("statistics", False),
            "files": files,
            "summary": summarize_collection(files),
        })
    grids = {dataset["name"]: _grid_signature(dataset["files"][0])
             for dataset in datasets if dataset["files"] and dataset["files"][0].get("status") == "ok"}
    pair_specs = config.get("paired_initializations", [])
    if selected:
        pair_specs = [spec for spec in pair_specs if set(spec["datasets"]).issubset(selected)]
    paired = _paired_initialization_report(datasets, pair_specs)
    return {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "config": str(source),
        "datasets": datasets,
        "compatibility": {
            "rainfall_temperature_grid_match": grids.get("rainfall_observation") == grids.get("temperature_observation")
            if "rainfall_observation" in grids and "temperature_observation" in grids else None,
            "rainfall_mask_grid_match": grids.get("rainfall_observation") == grids.get("vietnam_climate_mask")
            if "rainfall_observation" in grids and "vietnam_climate_mask" in grids else None,
            "paired_initializations": paired,
        },
    }


def has_failures(report: dict[str, Any]) -> bool:
    for dataset in report["datasets"]:
        if not dataset["files"] or dataset["summary"]["error_file_count"]:
            return True
        if dataset["summary"]["missing_expected_variable_file_counts"]:
            return True
        initialization = dataset["summary"].get("initialization_months")
        if initialization is not None and not initialization["continuous"]:
            return True
    if any(not pair["match"] for pair in report.get("compatibility", {}).get("paired_initializations", [])):
        return True
    return False


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="config/inventory.json", help="Inventory JSON configuration")
    parser.add_argument("--dataset", action="append", help="Dataset name to inspect; repeat as needed")
    parser.add_argument("--output", type=Path, help="Write JSON report here; stdout when omitted")
    parser.add_argument("--strict", action="store_true", help="Exit 1 for integrity, continuity or pairing failures")
    parser.add_argument("--log-level", default="INFO", choices=("DEBUG", "INFO", "WARNING", "ERROR"))
    args = parser.parse_args()
    logging.basicConfig(level=getattr(logging, args.log_level), format="%(levelname)s %(message)s")
    report = build_inventory(args.config, set(args.dataset) if args.dataset else None)
    rendered = json.dumps(report, indent=2, ensure_ascii=False)
    if args.output:
        output = project_path(args.output)
        output.parent.mkdir(parents=True, exist_ok=True)
        output.write_text(rendered + "\n", encoding="utf-8")
        print(output)
    else:
        print(rendered)
    if args.strict and has_failures(report):
        raise SystemExit(1)


if __name__ == "__main__":
    main()