"""Recover one ECMWF instantaneous file with resumable adaptive requests.

This is an exceptional-path tool for a month whose combined CDS/MARS request
hits a damaged tape.  It requests one variable at a time, starts with bounded
lead-time chunks, and bisects only chunks that return a damaged-tape error.
The normal raw file is promoted only after every expected lead and member has
been merged and validated.
"""
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, Callable

import numpy as np
import xarray as xr

from .config import PROJECT_ROOT
from .download_ecmwf_asia import load_domain
from .download_ecmwf_baseline import (
    _lead_hours,
    lead_hours_for_group,
    parse_init,
    validate_group_file,
)


LOG = logging.getLogger(__name__)


class ExternalDataUnavailableError(RuntimeError):
    """Raised when ECMWF confirms that a required single lead is unavailable."""


def split_chunks(values: list[str], size: int) -> list[list[str]]:
    if size < 1:
        raise ValueError("chunk size must be at least one")
    return [values[index:index + size] for index in range(0, len(values), size)]


def damaged_tape_error(exc: BaseException) -> bool:
    message = str(exc).lower()
    return "damaged tape" in message or ("accesserror" in message and "tape" in message)


def _atomic_json(path: Path, payload: dict[str, Any]) -> 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 _chunk_key(variable: str, leads: list[str]) -> str:
    return f"{variable}:{leads[0]}-{leads[-1]}:{len(leads)}"


def _chunk_path(work_directory: Path, variable: str, leads: list[str]) -> Path:
    return work_directory / variable / f"{variable}_{leads[0]}_{leads[-1]}_{len(leads)}.nc"


def reset_blocked_attempt_cycle(manifest: dict[str, Any],
                                previous: dict[str, Any]) -> list[str]:
    """Start a bounded new retry cycle while retaining lifetime provenance."""
    if previous.get("status") != "blocked_external":
        return []
    reset_utc = datetime.now(timezone.utc).isoformat()
    reset_keys = []
    for key, chunk in manifest.get("chunks", {}).items():
        leads = chunk.get("leads", [])
        attempts = int(chunk.get("attempt_count", 0))
        if (
            chunk.get("status") == "error"
            and len(leads) == 1
            and damaged_tape_error(RuntimeError(chunk.get("error", "")))
            and attempts > 0
        ):
            chunk["total_attempt_count"] = int(
                chunk.get("total_attempt_count", attempts)
            )
            chunk["attempt_count"] = 0
            chunk["attempt_cycle_reset_utc"] = reset_utc
            reset_keys.append(key)
    cycles = list(previous.get("recovery_cycles", []))
    cycles.append({
        "previous_status": previous.get("status"),
        "previous_updated_utc": previous.get("updated_utc"),
        "blocked_reason": previous.get("blocked_reason"),
        "reset_utc": reset_utc,
        "reset_chunks": reset_keys,
    })
    manifest["recovery_cycles"] = cycles
    return reset_keys


def _request(init: str, domain: dict[str, Any], cds_variable: str,
             leads: list[str]) -> dict[str, Any]:
    year, month = parse_init(init)
    return {
        "originating_centre": domain["originating_centre"],
        "system": domain["system"],
        "variable": [cds_variable],
        "year": str(year),
        "month": f"{month:02d}",
        "day": "01",
        "leadtime_hour": leads,
        "area": domain["area"],
        "data_format": "netcdf",
    }


def _variable_spec(group_spec: dict[str, Any], cds_variable: str,
                   netcdf_variable: str) -> dict[str, Any]:
    return {
        "variables": [cds_variable],
        "expected_netcdf_variables": [netcdf_variable],
        "expected_members": group_spec["expected_members"],
        "allow_missing_initial_step": False,
    }


def retrieve_adaptive(
    client: Any,
    dataset: str,
    init: str,
    domain: dict[str, Any],
    variable: str,
    cds_variable: str,
    leads: list[str],
    work_directory: Path,
    manifest: dict[str, Any],
    manifest_path: Path,
    validator: Callable[..., dict[str, Any]] = validate_group_file,
    max_single_step_attempts: int = 3,
) -> list[Path]:
    """Retrieve one chunk, bisecting damaged-tape failures until successful."""
    key = _chunk_key(variable, leads)
    target = _chunk_path(work_directory, variable, leads)
    partial = target.with_suffix(target.suffix + ".part")
    spec = _variable_spec(
        domain["download_groups"]["instantaneous"], cds_variable, variable
    )
    previous = manifest["chunks"].get(key, {})
    previous_attempts = int(previous.get("attempt_count", 0))
    previous_total_attempts = int(
        previous.get("total_attempt_count", previous_attempts)
    )
    if (
        len(leads) == 1
        and previous.get("status") == "error"
        and damaged_tape_error(RuntimeError(previous.get("error", "")))
        and previous_attempts >= max_single_step_attempts
    ):
        raise ExternalDataUnavailableError(
            f"{variable} lead {leads[0]} remains unavailable on a damaged "
            f"ECMWF tape after {previous_attempts} attempts"
        )
    if previous.get("status") == "split" and len(leads) > 1:
        middle = len(leads) // 2
        return (
            retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[:middle], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
            + retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[middle:], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
        )
    if target.exists():
        report = validator(target, init, spec, leads)
        manifest["chunks"][key] = {
            "status": "validated_existing", "path": str(target),
            "leads": leads, "validation": report,
        }
        _atomic_json(manifest_path, manifest)
        return [target]
    single_step_variables = manifest.setdefault("single_step_variables", [])
    if variable in single_step_variables and len(leads) > 1:
        middle = len(leads) // 2
        manifest["chunks"][key] = {
            "status": "split", "path": str(target), "leads": leads,
            "reason": "single-step mode after a two-lead damaged-tape failure",
            "children": [
                _chunk_key(variable, leads[:middle]),
                _chunk_key(variable, leads[middle:]),
            ],
        }
        _atomic_json(manifest_path, manifest)
        return (
            retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[:middle], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
            + retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[middle:], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
        )
    if partial.exists():
        failed = partial.with_suffix(partial.suffix + ".failed")
        if failed.exists():
            failed = partial.with_suffix(
                partial.suffix + "." + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ".failed"
            )
        os.replace(partial, failed)
    target.parent.mkdir(parents=True, exist_ok=True)
    request = _request(init, domain, cds_variable, leads)
    manifest["chunks"][key] = {
        "status": "requesting", "path": str(target), "leads": leads,
        "request": request, "updated_utc": datetime.now(timezone.utc).isoformat(),
        "attempt_count": previous_attempts + 1,
        "total_attempt_count": previous_total_attempts + 1,
    }
    _atomic_json(manifest_path, manifest)
    LOG.info("REQUEST variable=%s leads=%s..%s count=%s", variable,
             leads[0], leads[-1], len(leads))
    try:
        client.retrieve(dataset, request, str(partial))
        report = validator(partial, init, spec, leads)
        os.replace(partial, target)
        manifest["chunks"][key].update(
            status="downloaded", validation=report,
            updated_utc=datetime.now(timezone.utc).isoformat(),
        )
        _atomic_json(manifest_path, manifest)
        return [target]
    except Exception as exc:
        manifest["chunks"][key].update(
            status="error", error=f"{type(exc).__name__}: {exc}",
            updated_utc=datetime.now(timezone.utc).isoformat(),
        )
        _atomic_json(manifest_path, manifest)
        if not damaged_tape_error(exc):
            raise
        if len(leads) == 1:
            attempts = manifest["chunks"][key]["attempt_count"]
            if attempts >= max_single_step_attempts:
                raise ExternalDataUnavailableError(
                    f"{variable} lead {leads[0]} remains unavailable on a damaged "
                    f"ECMWF tape after {attempts} attempts"
                ) from exc
            raise
        middle = len(leads) // 2
        if len(leads) == 2:
            single_step_variables = manifest.setdefault("single_step_variables", [])
            if variable not in single_step_variables:
                single_step_variables.append(variable)
        manifest["chunks"][key]["status"] = "split"
        manifest["chunks"][key]["children"] = [
            _chunk_key(variable, leads[:middle]),
            _chunk_key(variable, leads[middle:]),
        ]
        _atomic_json(manifest_path, manifest)
        LOG.warning("SPLIT damaged-tape chunk variable=%s leads=%s..%s count=%s",
                    variable, leads[0], leads[-1], len(leads))
        return (
            retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[:middle], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
            + retrieve_adaptive(
                client, dataset, init, domain, variable, cds_variable,
                leads[middle:], work_directory, manifest, manifest_path, validator,
                max_single_step_attempts,
            )
        )


def _load_variable_chunks(paths: list[Path], variable: str,
                          expected_leads: list[str]) -> xr.DataArray:
    arrays = []
    for path in paths:
        with xr.open_dataset(path, decode_times=True) as dataset:
            if variable not in dataset:
                raise ValueError(f"{path} does not contain {variable}")
            arrays.append(dataset[variable].load())
    combined = xr.concat(arrays, dim="forecast_period").sortby("forecast_period")
    actual = _lead_hours(combined.forecast_period.values)
    expected = [int(value) for value in expected_leads]
    if actual != expected:
        duplicates = sorted({value for value in actual if actual.count(value) > 1})
        missing = sorted(set(expected) - set(actual))
        extra = sorted(set(actual) - set(expected))
        raise ValueError(
            f"Merged {variable} lead mismatch: missing={missing}, extra={extra}, "
            f"duplicates={duplicates}"
        )
    return combined


def merge_chunks(
    chunk_paths: dict[str, list[Path]],
    expected_leads: list[str],
    target: Path,
    init: str,
    group_spec: dict[str, Any],
    validator: Callable[..., dict[str, Any]] = validate_group_file,
) -> dict[str, Any]:
    """Merge validated variable chunks and atomically promote the raw file."""
    if target.exists():
        return validator(target, init, group_spec, ["0", *expected_leads])
    fields = [
        _load_variable_chunks(chunk_paths[name], name, expected_leads)
        for name in group_spec["expected_netcdf_variables"]
    ]
    merged = xr.merge(fields, compat="equals", join="exact")
    partial = target.with_suffix(target.suffix + ".part")
    if partial.exists():
        failed = partial.with_suffix(partial.suffix + ".failed")
        if failed.exists():
            failed = partial.with_suffix(
                partial.suffix + "." + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + ".failed"
            )
        os.replace(partial, failed)
    target.parent.mkdir(parents=True, exist_ok=True)
    encoding = {
        name: {"zlib": True, "complevel": 1, "shuffle": True}
        for name in merged.data_vars
    }
    merged.attrs.update({
        "recovery_method": "per-variable adaptive lead chunking after ECMWF damaged-tape error",
        "recovery_initial_month": init,
        "created_utc": datetime.now(timezone.utc).isoformat(),
    })
    merged.to_netcdf(partial, encoding=encoding)
    report = validator(partial, init, group_spec, ["0", *expected_leads])
    os.replace(partial, target)
    return report


def run_recovery(init: str, domain: dict[str, Any], work_directory: Path,
                 target: Path, chunk_leads: int = 112, client: Any | None = None,
                 dry_run: bool = False,
                 max_single_step_attempts: int = 3,
                 reset_blocked_attempts: bool = False) -> tuple[Path, dict[str, Any]]:
    parse_init(init)
    if max_single_step_attempts < 1:
        raise ValueError("max_single_step_attempts must be at least one")
    spec = domain["download_groups"]["instantaneous"]
    requested = lead_hours_for_group(init, spec)
    leads = requested[1:] if requested and requested[0] == "0" else requested
    variables = dict(zip(spec["expected_netcdf_variables"], spec["variables"]))
    initial_chunks = split_chunks(leads, chunk_leads)
    manifest_path = work_directory / "recovery_manifest.json"
    manifest = {
        "initial_month": init,
        "dataset": domain["dataset"],
        "area": domain["area"],
        "target": str(target),
        "expected_leads": leads,
        "variables": variables,
        "initial_chunk_leads": chunk_leads,
        "chunks": {},
        "single_step_variables": [],
        "status": "dry_run" if dry_run else "running",
        "updated_utc": datetime.now(timezone.utc).isoformat(),
    }
    if manifest_path.exists() and not dry_run:
        previous = json.loads(manifest_path.read_text(encoding="utf-8"))
        if previous.get("initial_month") != init or previous.get("expected_leads") != leads:
            raise ValueError("Existing recovery manifest does not match this run")
        manifest["chunks"] = previous.get("chunks", {})
        manifest["single_step_variables"] = previous.get("single_step_variables", [])
        manifest["recovery_cycles"] = previous.get("recovery_cycles", [])
        if reset_blocked_attempts:
            reset_blocked_attempt_cycle(manifest, previous)
    if dry_run:
        manifest["planned_initial_requests"] = len(initial_chunks) * len(variables)
        manifest["lead_count"] = len(leads)
        return manifest_path, manifest
    if client is None:
        try:
            import cdsapi
        except ImportError as exc:
            raise RuntimeError("cdsapi is required") from exc
        client = cdsapi.Client()
    _atomic_json(manifest_path, manifest)
    chunk_paths: dict[str, list[Path]] = {}
    try:
        for variable, cds_variable in variables.items():
            paths = []
            for chunk in initial_chunks:
                paths.extend(retrieve_adaptive(
                    client, domain["dataset"], init, domain, variable,
                    cds_variable, chunk, work_directory, manifest, manifest_path,
                    max_single_step_attempts=max_single_step_attempts,
                ))
            chunk_paths[variable] = paths
        report = merge_chunks(chunk_paths, leads, target, init, spec)
    except ExternalDataUnavailableError as exc:
        manifest["status"] = "blocked_external"
        manifest["blocked_reason"] = str(exc)
        manifest["updated_utc"] = datetime.now(timezone.utc).isoformat()
        _atomic_json(manifest_path, manifest)
        raise
    except Exception:
        manifest["status"] = "error"
        manifest["updated_utc"] = datetime.now(timezone.utc).isoformat()
        _atomic_json(manifest_path, manifest)
        raise
    manifest["status"] = "complete"
    manifest["validation"] = report
    manifest["updated_utc"] = datetime.now(timezone.utc).isoformat()
    _atomic_json(manifest_path, manifest)
    return manifest_path, manifest


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--init", required=True)
    parser.add_argument("--domain-config", default="config/ecmwf_domain.json")
    parser.add_argument("--work-directory", type=Path)
    parser.add_argument("--target", type=Path)
    parser.add_argument("--chunk-leads", type=int, default=112)
    parser.add_argument("--max-single-step-attempts", type=int, default=3,
                        help="Stop after this many damaged-tape failures for one lead")
    parser.add_argument("--reset-blocked-attempts", action="store_true",
                        help="Start a new bounded retry cycle after cooldown")
    parser.add_argument("--dry-run", action="store_true")
    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")
    domain = load_domain(args.domain_config)
    work = args.work_directory or Path(f"data/ecmwf/recovery/{args.init}-instantaneous")
    target = args.target or Path(domain["output_directory"]) / f"ecmwf_seas51_{args.init}_instantaneous.nc"
    if not work.is_absolute():
        work = PROJECT_ROOT / work
    if not target.is_absolute():
        target = PROJECT_ROOT / target
    manifest_path, manifest = run_recovery(
        args.init, domain, work, target, args.chunk_leads,
        dry_run=args.dry_run,
        max_single_step_attempts=args.max_single_step_attempts,
        reset_blocked_attempts=args.reset_blocked_attempts,
    )
    print(json.dumps({"manifest": str(manifest_path), "status": manifest["status"],
                      "target": str(target),
                      "planned_initial_requests": manifest.get("planned_initial_requests")},
                     indent=2))


if __name__ == "__main__":
    main()