#!/usr/bin/env python3
"""Download an ECMWF S2S *realtime* control forecast at native 1.5 deg (ECDS).

Mirror of download_ec15.py but for the `s2s-forecasts` dataset: one real init
date (no model-version / hindcast keys). Reuses the exact variable / level /
leadtime / area definitions of ec15_groups.py so the resulting GRIB is a
1.5 deg twin of what the reforecast download produced, ready to be turned into
EC_0.25_SD for a single operational init.

Usage:
  download_ec15_rt.py --init 2026-08-29 [--only srf_daily txtm plev ...]
"""
import argparse
import datetime as dt
import os
import sys

sys.path.insert(0, os.path.join(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_rt")
DATASET_RT = "s2s-forecasts"


def build_request_rt(group, init):
    """One realtime ECDS request for `group` at init date `init` (dt.date)."""
    g = dict(G.GROUPS[group])                       # copy: variable/level/leadtime/area
    req = {
        "origin": "ecmwf",
        "year": "%04d" % init.year,
        "month": "%02d" % init.month,
        "day": "%02d" % init.day,
        "time": "00:00",
        "forecast_type": "control_forecast",
        "data_format": "grib",
        "grid": G.GRID,
    }
    req.update(g)
    return req


def fetch_one(group, init, retries=3):
    outdir = os.path.join(RAWDIR, group)
    os.makedirs(outdir, exist_ok=True)
    target = os.path.join(outdir, "%s_%s.grib" % (group, init.isoformat()))
    tmp = target + ".part"
    if os.path.exists(target) and os.path.getsize(target) > 0:
        print("  skip (exists): %s" % os.path.basename(target), flush=True)
        return True
    req = build_request_rt(group, init)
    for attempt in range(1, retries + 1):
        try:
            if os.path.exists(tmp):
                os.remove(tmp)
            cdsapi.Client(url=G.API_URL, quiet=True, progress=False,
                          retry_max=20, sleep_max=120).retrieve(
                DATASET_RT, req).download(tmp)
            if not os.path.exists(tmp) or os.path.getsize(tmp) == 0:
                raise RuntimeError("empty download")
            os.replace(tmp, target)
            print("  ok: %s (%.1f MB)"
                  % (os.path.basename(target), os.path.getsize(target) / 1e6),
                  flush=True)
            return True
        except Exception as exc:                       # noqa: BLE001
            print("  [%s %s] attempt %d/%d failed: %s"
                  % (group, init, attempt, retries, exc), flush=True)
            if attempt == retries:
                return False
    return False


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", required=True, help="init date YYYY-MM-DD")
    ap.add_argument("--only", nargs="*", default=None,
                    help="subset of groups; default = all in ec15_groups")
    args = ap.parse_args()
    init = dt.date.fromisoformat(args.init)
    groups = [g for g in sorted(G.GROUPS) if args.only is None or g in args.only]
    print("[rt] realtime s2s-forecasts | init %s | grid %s | %d groups: %s"
          % (init, G.GRID, len(groups), ", ".join(groups)), flush=True)
    ok = 0
    for g in groups:
        print("[rt] %s ->" % g, flush=True)
        ok += bool(fetch_one(g, init))
    print("[rt] %d/%d groups done" % (ok, len(groups)), flush=True)
    return 0 if ok == len(groups) else 1


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