"""Run the approved small CDS pilot with variables split by GRIB semantics.

Each group is downloaded to a temporary ``.part`` file, validated as NetCDF,
then atomically promoted.  Failures retain the temporary file for diagnosis
and are recorded in a JSON manifest.  This command is not the baseline
downloader.
"""
from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import numpy as np
import xarray as xr

from .config import PROJECT_ROOT, project_path
from .download_ecmwf import DATASET
from .download_ecmwf_asia import load_domain
from .inventory import inventory_file


LOG = logging.getLogger(__name__)


def load_pilot_config(path: str | Path = "config/ecmwf_pilot.json") -> dict[str, Any]:
    source = project_path(path)
    with source.open(encoding="utf-8") as handle:
        config = json.load(handle)
    config["_source"] = str(source)
    return config


def request_for_group(init: str, domain: dict[str, Any], pilot: dict[str, Any], group: str) -> dict[str, Any]:
    try:
        year_text, month_text = init.split("-")
        year, month = int(year_text), int(month_text)
    except (TypeError, ValueError) as exc:
        raise ValueError("init must be formatted YYYY-MM") from exc
    if len(year_text) != 4 or len(month_text) != 2 or not 1 <= month <= 12:
        raise ValueError("init must be formatted YYYY-MM with a valid month")
    if group not in pilot["groups"]:
        raise KeyError(f"Unknown pilot group {group!r}; choose from {sorted(pilot['groups'])}")
    specification = pilot["groups"][group]
    leadtime_hour = specification.get("leadtime_hour")
    if leadtime_hour is None:
        leadtime_hour = pilot.get("leadtime_hour")
    if not leadtime_hour:
        raise ValueError(f"Pilot group {group!r} must define non-empty leadtime_hour")
    return {
        "originating_centre": domain["originating_centre"],
        "system": domain["system"],
        "variable": specification["variables"],
        "year": str(year),
        "month": f"{month:02d}",
        "day": "01",
        "leadtime_hour": [str(value) for value in leadtime_hour],
        "area": domain["area"],
        "data_format": "netcdf",
    }


def validate_pilot_file(
    path: Path,
    expected_variables: list[str],
    expected_leadtime_hour: list[str] | None = None,
    expected_soil_layers: int | None = None,
) -> dict[str, Any]:
    report = inventory_file(path, expected_variables=expected_variables, statistics=False)
    if report["status"] != "ok":
        raise ValueError(f"Downloaded file is not readable NetCDF: {report.get('error', 'unknown error')}")
    missing = report["missing_expected_variables"]
    if missing:
        raise ValueError(f"Downloaded file is missing expected variable(s): {missing}")
    with xr.open_dataset(path, decode_times=True) as ds:
        member_name = next((name for name in ("number", "member", "realization") if name in ds.sizes), None)
        if member_name is None:
            raise ValueError("Downloaded file has no ensemble-member dimension")
        if int(ds.sizes[member_name]) != 25:
            raise ValueError(f"Expected 25 ensemble members; found {ds.sizes[member_name]}")
        if expected_leadtime_hour is not None:
            lead_name = next((name for name in ("forecast_period", "step", "lead") if name in ds), None)
            if lead_name is None:
                raise ValueError("Downloaded file has no forecast lead coordinate")
            lead = ds[lead_name]
            if np.issubdtype(lead.dtype, np.timedelta64):
                actual = [str(int(value / np.timedelta64(1, "h"))) for value in lead.values]
            else:
                actual = [str(int(value)) for value in lead.values]
            if actual != [str(value) for value in expected_leadtime_hour]:
                raise ValueError(f"Expected lead hours {expected_leadtime_hour}; found {actual}")
        if expected_soil_layers is not None:
            soil_name = next(
                (name for name in ("soilLayer", "soil_layer", "depthBelowLandLayer") if name in ds.sizes),
                None,
            )
            if soil_name is None:
                raise ValueError("Downloaded soil-moisture file has no soil-layer dimension")
            if int(ds.sizes[soil_name]) != expected_soil_layers:
                raise ValueError(
                    f"Expected {expected_soil_layers} soil layers; "
                    f"found {ds.sizes[soil_name]} on {soil_name}"
                )
    return report


def _target_path(output_directory: Path, init: str, group: str) -> Path:
    return output_directory / f"ecmwf_seas51_pilot_{init}_{group}.nc"


def retrieve_group(
    client: Any,
    init: str,
    group: str,
    domain: dict[str, Any],
    pilot: dict[str, Any],
    output_directory: Path,
    overwrite: bool = False,
    retry_partial: bool = False,
) -> dict[str, Any]:
    specification = pilot["groups"][group]
    request = request_for_group(init, domain, pilot, group)
    target = _target_path(output_directory, init, group)
    partial = target.with_suffix(target.suffix + ".part")
    if target.exists() and not overwrite:
        report = validate_pilot_file(
            target,
            specification["expected_netcdf_variables"],
            specification.get("expected_leadtime_hour"),
            specification.get("expected_soil_layers"),
        )
        return {"group": group, "status": "validated_existing", "path": str(target), "request": request, "inventory": report}
    if partial.exists() and retry_partial and not overwrite:
        partial.unlink()
    if partial.exists() and not overwrite:
        raise FileExistsError(f"Partial pilot output exists; inspect it or use --overwrite: {partial}")
    if overwrite:
        target.unlink(missing_ok=True)
        partial.unlink(missing_ok=True)
    output_directory.mkdir(parents=True, exist_ok=True)
    LOG.info("Submitting pilot group=%s target=%s", group, target)
    client.retrieve(DATASET, request, str(partial))
    report = validate_pilot_file(
        partial,
        specification["expected_netcdf_variables"],
        specification.get("expected_leadtime_hour"),
        specification.get("expected_soil_layers"),
    )
    os.replace(partial, target)
    return {"group": group, "status": "downloaded", "path": str(target), "request": request, "inventory": report}


def _write_manifest(path: Path, manifest: dict[str, Any]) -> None:
    temporary = path.with_suffix(path.suffix + ".part")
    temporary.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    os.replace(temporary, path)


def run_pilot(
    init: str,
    groups: list[str],
    pilot: dict[str, Any],
    domain: dict[str, Any],
    output_directory: Path,
    overwrite: bool = False,
    retry_partial: bool = False,
    client: Any | None = None,
) -> tuple[Path, dict[str, Any]]:
    if client is None:
        try:
            import cdsapi
        except ImportError as exc:
            raise RuntimeError("cdsapi is required in the active Python environment") from exc
        client = cdsapi.Client()
    manifest: dict[str, Any] = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "dataset": DATASET,
        "initial_month": init,
        "domain": domain["area"],
        "leadtime_hour_by_group": {
            group: request_for_group(init, domain, pilot, group)["leadtime_hour"]
            for group in groups
        },
        "pilot_config": pilot.get("_source"),
        "results": [],
    }
    manifest_path = output_directory / f"ecmwf_seas51_pilot_{init}_manifest.json"
    output_directory.mkdir(parents=True, exist_ok=True)
    _write_manifest(manifest_path, manifest)
    for group in groups:
        try:
            result = retrieve_group(
                client, init, group, domain, pilot, output_directory,
                overwrite, retry_partial,
            )
        except Exception as exc:
            LOG.exception("Pilot group failed: %s", group)
            result = {
                "group": group,
                "status": "error",
                "error": f"{type(exc).__name__}: {exc}",
                "request": request_for_group(init, domain, pilot, group),
            }
        manifest["results"].append(result)
        # Persist progress after every request so an interrupted CDS process
        # still leaves a diagnostic checkpoint for the completed group(s).
        manifest["updated_at"] = datetime.now(timezone.utc).isoformat()
        _write_manifest(manifest_path, manifest)
    return manifest_path, manifest


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--init", help="Initial month YYYY-MM; defaults to pilot config")
    parser.add_argument("--group", action="append", help="Group to download; repeat as needed; defaults to all groups")
    parser.add_argument("--pilot-config", default="config/ecmwf_pilot.json")
    parser.add_argument("--domain-config", default="config/ecmwf_domain.json")
    parser.add_argument("--output-directory", type=Path, help="Override configured pilot output directory")
    parser.add_argument("--overwrite", action="store_true", help="Replace existing pilot/partial files")
    parser.add_argument(
        "--retry-partial", action="store_true",
        help="Remove and retry only an incomplete .part pilot file",
    )
    parser.add_argument("--dry-run", action="store_true", help="Print the three requests without contacting CDS")
    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="%(asctime)s %(levelname)s %(message)s")
    pilot = load_pilot_config(args.pilot_config)
    domain = load_domain(args.domain_config)
    init = args.init or pilot["initial_month"]
    groups = args.group or list(pilot["groups"])
    unknown = sorted(set(groups) - set(pilot["groups"]))
    if unknown:
        parser.error(f"unknown group(s): {unknown}; choices: {sorted(pilot['groups'])}")
    output_directory = project_path(args.output_directory or pilot["output_directory"])
    if args.dry_run:
        print(json.dumps({group: request_for_group(init, domain, pilot, group) for group in groups}, indent=2))
        return
    manifest_path, manifest = run_pilot(
        init, groups, pilot, domain, output_directory,
        args.overwrite, args.retry_partial,
    )
    print(manifest_path)
    if any(result["status"] == "error" for result in manifest["results"]):
        raise SystemExit(1)


if __name__ == "__main__":
    main()