#!/usr/bin/env python3
"""Download ECMWF S2S reforecasts (native 1.5 deg) from ECDS, one group at a time.

One request = one model-version date x all 20 hindcast years, which cuts the
request count from ~2100 to 104 per group. Raw GRIB is kept; conversion to
NetCDF is done by convert_ec15.py so a failed conversion never forces a re-download.

Usage:
  download_ec15.py <group> [--workers N] [--dates mm-dd,mm-dd] [--dry-run]
"""
import argparse
import os
import sys
import time
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ec15_groups as G

import cdsapi

ROOT = "/home/tanpv/Progs/claude/HW_EC_10"
RAWDIR = os.path.join(ROOT, "data", "grib")

_print_lock = threading.Lock()


def log(msg):
    with _print_lock:
        print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}", flush=True)


def fetch_one(group, mvdate, retries=3):
    """Download one (group, model-version date). Returns (mvdate, status, bytes)."""
    outdir = os.path.join(RAWDIR, group)
    os.makedirs(outdir, exist_ok=True)
    target = os.path.join(outdir, f"{group}_{mvdate}.grib")
    tmp = target + ".part"

    if os.path.exists(target) and os.path.getsize(target) > 0:
        return (mvdate, "skip", os.path.getsize(target))

    mm, dd = mvdate.split("-")
    req = G.build_request(group, mm, dd)

    for attempt in range(1, retries + 1):
        try:
            if os.path.exists(tmp):
                os.remove(tmp)
            client = cdsapi.Client(url=G.API_URL, quiet=True, progress=False,
                                   retry_max=20, sleep_max=120)
            client.retrieve(G.DATASET, req).download(tmp)
            if not os.path.exists(tmp) or os.path.getsize(tmp) == 0:
                raise RuntimeError("empty download")
            os.replace(tmp, target)
            return (mvdate, "ok", os.path.getsize(target))
        except Exception as exc:                       # noqa: BLE001 - report and retry
            log(f"  [{group} {mvdate}] attempt {attempt}/{retries} failed: {exc}")
            if attempt == retries:
                return (mvdate, f"FAIL: {exc}", 0)
            time.sleep(30 * attempt)
    return (mvdate, "FAIL: unreachable", 0)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("group", choices=sorted(G.GROUPS))
    ap.add_argument("--workers", type=int, default=3,
                    help="concurrent ECDS requests (ECDS limits per-user concurrency)")
    ap.add_argument("--dates", default=None, help="comma-separated mm-dd subset")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    dates = args.dates.split(",") if args.dates else G.load_mvdates()
    log(f"group={args.group} ndates={len(dates)} workers={args.workers}")

    if args.dry_run:
        req = G.build_request(args.group, *dates[0].split("-"))
        for k, v in req.items():
            log(f"  {k}: {v if not isinstance(v, list) or len(v) < 8 else str(v[:5])+f'... ({len(v)})'}")
        return 0

    t0 = time.time()
    nok = nskip = nfail = 0
    total_bytes = 0
    with ThreadPoolExecutor(max_workers=args.workers) as pool:
        futs = {pool.submit(fetch_one, args.group, d): d for d in dates}
        for i, fut in enumerate(as_completed(futs), 1):
            mvdate, status, nbytes = fut.result()
            total_bytes += nbytes
            if status == "ok":
                nok += 1
            elif status == "skip":
                nskip += 1
            else:
                nfail += 1
            log(f"[{i}/{len(dates)}] {args.group} {mvdate}: {status} "
                f"({nbytes/1e6:.1f} MB)")

    log(f"DONE group={args.group} ok={nok} skip={nskip} fail={nfail} "
        f"total={total_bytes/1e9:.2f} GB elapsed={(time.time()-t0)/60:.1f} min")
    return 1 if nfail else 0


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