"""Safe, resumable ECMWF SEAS5.1 baseline downloader."""
from __future__ import annotations

import argparse
import json
import os
import shutil
from contextlib import ExitStack
from datetime import datetime, timezone
from pathlib import Path

from .config import PROJECT_ROOT
from .download_ecmwf import lead_hours_through_lt7_boundary
from .download_ecmwf_asia import load_domain


def parse_init(init: str) -> tuple[int, int]:
    try:
        parsed = datetime.strptime(init, "%Y-%m")
    except ValueError as exc:
        raise ValueError(f"Invalid initial month {init!r}; expected YYYY-MM") from exc
    if parsed.strftime("%Y-%m") != init:
        raise ValueError(f"Invalid initial month {init!r}; expected YYYY-MM")
    return parsed.year, parsed.month


def configured_groups(domain: dict, requested: list[str] | None = None) -> list[str]:
    specs = domain["download_groups"]
    if requested:
        unknown = sorted(set(requested) - set(specs))
        if unknown:
            raise ValueError(f"Unknown download group(s): {', '.join(unknown)}")
        return list(dict.fromkeys(requested))
    return [name for name, spec in specs.items()
            if spec.get("enabled_by_default", True)]


def lead_hours_for_group(init: str, spec: dict) -> list[str]:
    year, month = parse_init(init)
    boundary = int(lead_hours_through_lt7_boundary(year, month)[-1])
    start, interval = int(spec["start_hour"]), int(spec["step_hours"])
    if interval <= 0 or start < 0 or (boundary - start) % interval:
        raise ValueError(
            f"Steps do not land on LT=7: start={start}, interval={interval}, "
            f"boundary={boundary}"
        )
    return [str(hour) for hour in range(start, boundary + 1, interval)]


def request_for_group(init: str, domain: dict, group: str) -> dict:
    year, month = parse_init(init)
    spec = domain["download_groups"][group]
    return {
        "originating_centre": domain["originating_centre"],
        "system": domain["system"],
        "variable": spec["variables"],
        "year": str(year),
        "month": f"{month:02d}",
        "day": "01",
        "leadtime_hour": lead_hours_for_group(init, spec),
        "area": domain["area"],
        "data_format": "netcdf",
        # Make the delivery mode explicit.  Besides guaranteeing a single
        # NetCDF target, this avoids reusing legacy CDS cache entries whose
        # object-store download URL may already have expired.
        "download_format": "unarchived",
    }


def _lead_hours(values) -> list[int]:
    import numpy as np
    array = np.asarray(values)
    if np.issubdtype(array.dtype, np.timedelta64):
        return (array / np.timedelta64(1, "h")).astype("int64").tolist()
    return array.astype("int64").tolist()


def validate_group_file(path: Path, init: str, spec: dict,
                        requested_leads: list[str]) -> dict:
    import numpy as np
    import xarray as xr
    with xr.open_dataset(path, decode_times=True) as ds:
        missing = sorted(set(spec["expected_netcdf_variables"]) - set(ds.data_vars))
        if missing:
            raise ValueError(f"Missing NetCDF variables: {', '.join(missing)}")
        member = next((x for x in ("number", "member", "realization")
                       if x in ds.dims), None)
        count = None if member is None else ds.sizes[member]
        if count != int(spec["expected_members"]):
            raise ValueError(
                f"Expected {spec['expected_members']} ensemble members; found {count}"
            )
        expected_soil_layers = spec.get("expected_soil_layers")
        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("No soil-layer dimension found")
            if int(ds.sizes[soil_name]) != int(expected_soil_layers):
                raise ValueError(
                    f"Expected {expected_soil_layers} soil layers; "
                    f"found {ds.sizes[soil_name]} on {soil_name}"
                )
        lead = next((x for x in ("forecast_period", "step", "leadtime_hour")
                     if x in ds), None)
        if lead is None:
            raise ValueError("No forecast period coordinate found")
        actual = _lead_hours(ds[lead].values)
        expected = [int(x) for x in requested_leads]
        accepted = [expected]
        if spec.get("allow_missing_initial_step") and expected[0] == 0:
            accepted.append(expected[1:])
        if actual not in accepted:
            raise ValueError(
                f"Lead-time mismatch: expected lengths {[len(x) for x in accepted]}, "
                f"found {len(actual)} ({actual[0] if actual else None}.."
                f"{actual[-1] if actual else None})"
            )
        ref = next((x for x in ("forecast_reference_time", "time") if x in ds), None)
        if ref is None:
            raise ValueError("No forecast reference time coordinate found")
        ref_values = np.asarray(ds[ref].values).reshape(-1)
        expected_ref = np.datetime64(f"{init}-01T00:00:00")
        if (ref_values.size != 1 or
                ref_values[0].astype("datetime64[s]") != expected_ref):
            raise ValueError(
                f"Forecast reference time mismatch: {ref_values}; expected {expected_ref}"
            )
        report = {
            "variables": sorted(spec["expected_netcdf_variables"]),
            "members": int(count), "first_lead_hour": actual[0],
            "last_lead_hour": actual[-1], "lead_count": len(actual),
            "initial_step_absent": actual == expected[1:],
            "bytes": path.stat().st_size,
        }
        if expected_soil_layers is not None:
            report["soil_layers"] = int(expected_soil_layers)
            report["soil_layer_coordinate"] = soil_name
            variable = ds[spec["expected_netcdf_variables"][0]]
            report["units"] = variable.attrs.get("units")
            report["grib_step_type"] = variable.attrs.get("GRIB_stepType")
        return report


def _write_manifest(path: Path, payload: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    partial = path.with_suffix(path.suffix + ".part")
    partial.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n",
                       encoding="utf-8")
    os.replace(partial, path)


def _is_expired_object_store_error(exc: Exception) -> bool:
    """Return true only for the stale CDS object-store URL failure."""
    response = getattr(exc, "response", None)
    status = getattr(response, "status_code", None)
    url = getattr(response, "url", "") or ""
    message = str(exc)
    return (
        (status == 403 or "403" in message)
        and ("object-store" in url or "object-store" in message)
    )


def _retrieve_split_by_variable(
    client, dataset: str, init: str, spec: dict, request: dict,
    target: Path, retry_partial: bool, validator,
) -> dict:
    """Bypass one stale multi-variable cache entry with resumable pieces."""
    import xarray as xr

    source_variables = list(request["variable"])
    output_variables = list(spec["expected_netcdf_variables"])
    if len(source_variables) != len(output_variables):
        raise ValueError(
            "Cannot split request: source and NetCDF variable counts differ "
            f"({len(source_variables)} != {len(output_variables)})"
        )

    parts_dir = target.parent / f"{target.name}.split-parts"
    parts_dir.mkdir(parents=True, exist_ok=True)
    part_paths = []
    part_reports = {}
    for source_variable, output_variable in zip(source_variables, output_variables):
        part_target = parts_dir / f"{output_variable}.nc"
        part_partial = part_target.with_suffix(".nc.part")
        part_request = dict(request)
        part_request["variable"] = [source_variable]
        part_spec = dict(spec)
        part_spec["expected_netcdf_variables"] = [output_variable]

        if part_target.exists():
            report = validator(
                part_target, init, part_spec, part_request["leadtime_hour"]
            )
        else:
            if part_partial.exists() and retry_partial:
                part_partial.unlink()
            if part_partial.exists():
                raise FileExistsError(
                    f"Split partial retained at {part_partial}; inspect it or retry"
                )
            client.retrieve(dataset, part_request, str(part_partial))
            report = validator(
                part_partial, init, part_spec, part_request["leadtime_hour"]
            )
            os.replace(part_partial, part_target)
        part_paths.append(part_target)
        part_reports[output_variable] = report

    combined_partial = target.with_suffix(target.suffix + ".part")
    combined_partial.unlink(missing_ok=True)
    with ExitStack() as stack:
        datasets = [stack.enter_context(xr.open_dataset(path)) for path in part_paths]
        combined = xr.merge(datasets, compat="identical", join="exact")
        combined.to_netcdf(combined_partial)
    report = validator(
        combined_partial, init, spec, request["leadtime_hour"]
    )
    os.replace(combined_partial, target)
    shutil.rmtree(parts_dir)
    return {
        "status": "downloaded_split_fallback",
        "path": str(target),
        "validation": report,
        "split_validation": part_reports,
    }


def retrieve_group(client, dataset: str, init: str, spec: dict, request: dict,
                   target: Path, overwrite: bool = False, retry_partial: bool = False,
                   validator=validate_group_file) -> dict:
    partial = target.with_suffix(target.suffix + ".part")
    if target.exists() and not overwrite:
        report = validator(target, init, spec, request["leadtime_hour"])
        return {"status": "validated_existing", "path": str(target),
                "validation": report}
    if partial.exists() and retry_partial and not overwrite:
        partial.unlink()
    if partial.exists() and not overwrite:
        raise FileExistsError(
            f"Partial file retained at {partial}; inspect it or use --overwrite"
        )
    if overwrite:
        target.unlink(missing_ok=True)
        partial.unlink(missing_ok=True)
    target.parent.mkdir(parents=True, exist_ok=True)
    try:
        client.retrieve(dataset, request, str(partial))
    except Exception as exc:
        if not _is_expired_object_store_error(exc):
            raise
        partial.unlink(missing_ok=True)
        result = _retrieve_split_by_variable(
            client, dataset, init, spec, request, target, retry_partial, validator
        )
        result["fallback_reason"] = f"{type(exc).__name__}: {exc}"
        return result
    report = validator(partial, init, spec, request["leadtime_hour"])
    os.replace(partial, target)
    return {"status": "downloaded", "path": str(target), "validation": report}


def run(init: str, domain: dict, groups: list[str], output: Path,
        overwrite: bool = False, dry_run: bool = False, client=None,
        retry_partial: bool = False) -> tuple[dict, bool]:
    parse_init(init)
    manifest_path = output / f"ecmwf_seas51_{init}_manifest.json"
    manifest = {
        "initial_month": init, "dataset": domain["dataset"], "area": domain["area"],
        "groups": {}, "updated_utc": datetime.now(timezone.utc).isoformat(),
    }
    if not dry_run and manifest_path.exists():
        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    if not dry_run and client is None:
        try:
            import cdsapi
        except ImportError as exc:
            raise SystemExit("cdsapi is required.") from exc
        client = cdsapi.Client()
    success = True
    for group in groups:
        spec = domain["download_groups"][group]
        request = request_for_group(init, domain, group)
        target = output / f"ecmwf_seas51_{init}_{group}.nc"
        entry = {"target": str(target), "request": request,
                 "temporal_semantics": spec["temporal_semantics"]}
        if dry_run:
            entry["status"] = "dry_run"
        else:
            try:
                entry.update(retrieve_group(
                    client, domain["dataset"], init, spec, request, target, overwrite, retry_partial
                ))
            except Exception as exc:
                success = False
                entry.update(status="error", error=f"{type(exc).__name__}: {exc}")
        manifest["groups"][group] = entry
        manifest["updated_utc"] = datetime.now(timezone.utc).isoformat()
        if not dry_run:
            _write_manifest(manifest_path, manifest)
    return manifest, success


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--init", required=True)
    parser.add_argument("--group", action="append", dest="groups")
    parser.add_argument("--overwrite", action="store_true")
    parser.add_argument("--retry-partial", action="store_true",
                        help="Remove and retry only incomplete .part files")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--domain-config", default="config/ecmwf_domain.json")
    parser.add_argument("--output-directory")
    args = parser.parse_args()
    domain = load_domain(args.domain_config)
    groups = configured_groups(domain, args.groups)
    output = Path(args.output_directory or domain["output_directory"])
    if not output.is_absolute():
        output = PROJECT_ROOT / output
    manifest, success = run(
        args.init, domain, groups, output, args.overwrite, args.dry_run, retry_partial=args.retry_partial
    )
    print(json.dumps(manifest, indent=2, sort_keys=True))
    raise SystemExit(0 if success else 1)


if __name__ == "__main__":
    main()