#!/usr/bin/env python3
"""Convert downloaded S2S GRIB (one model-version date x 20 hindcast years)
into per-hindcast-year NetCDF, mirroring the layout of the existing 0.25 deg
archive so downstream code can be reused:

    data/nc/<group>/<hyear>/<group>_cf_<hyear>-<mm-dd>.nc

Split by GRIB key `hdate` first, then grib_to_netcdf each year separately.
Conversion is kept apart from download so a conversion bug never costs a re-download.

Usage: convert_ec15.py <group> [--jobs N] [--force]
"""
import argparse
import glob
import os
import re
import shutil
import subprocess
import sys
import tempfile
from concurrent.futures import ProcessPoolExecutor, as_completed

ROOT = "/home/tanpv/Progs/claude/HW_EC_10"
RAWDIR = os.path.join(ROOT, "data", "grib")
NCDIR = os.path.join(ROOT, "data", "nc")
ECCODES_BIN = "/home/hapt/miniconda3/bin"

GRIB_COPY = os.path.join(ECCODES_BIN, "grib_copy")
GRIB_TO_NC = os.path.join(ECCODES_BIN, "grib_to_netcdf")


def run(cmd):
    p = subprocess.run(cmd, capture_output=True, text=True)
    if p.returncode != 0:
        raise RuntimeError(f"{' '.join(cmd[:2])} failed: {p.stderr.strip()[:400]}")
    return p.stdout


def convert_one(args):
    """Split one GRIB by hdate and write one NetCDF per hindcast year."""
    group, gribfile, force = args
    mvdate = re.search(r"_(\d{2}-\d{2})\.grib$", gribfile).group(1)
    written, skipped = [], []
    tmpdir = tempfile.mkdtemp(prefix=f"conv_{group}_{mvdate}_", dir="/tmp")
    try:
        run([GRIB_COPY, gribfile, os.path.join(tmpdir, "h_[hdate].grib")])
        for part in sorted(glob.glob(os.path.join(tmpdir, "h_*.grib"))):
            hdate = re.search(r"h_(\d{8})\.grib$", part).group(1)
            hyear = hdate[:4]
            outdir = os.path.join(NCDIR, group, hyear)
            os.makedirs(outdir, exist_ok=True)
            out = os.path.join(outdir, f"{group}_cf_{hyear}-{mvdate}.nc")
            if os.path.exists(out) and os.path.getsize(out) > 0 and not force:
                skipped.append(hyear)
                continue
            run([GRIB_TO_NC, "-o", out, part])
            written.append(hyear)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
    return (mvdate, written, skipped)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("group")
    ap.add_argument("--jobs", type=int, default=8)
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    gribs = sorted(glob.glob(os.path.join(RAWDIR, args.group, f"{args.group}_*.grib")))
    if not gribs:
        print(f"No GRIB found for group {args.group} under {RAWDIR}", file=sys.stderr)
        return 1
    print(f"group={args.group} nfiles={len(gribs)} jobs={args.jobs}", flush=True)

    nw = ns = nfail = 0
    tasks = [(args.group, g, args.force) for g in gribs]
    with ProcessPoolExecutor(max_workers=args.jobs) as pool:
        futs = {pool.submit(convert_one, t): t[1] for t in tasks}
        for i, fut in enumerate(as_completed(futs), 1):
            try:
                mvdate, written, skipped = fut.result()
                nw += len(written)
                ns += len(skipped)
                print(f"[{i}/{len(gribs)}] {mvdate}: {len(written)} written, "
                      f"{len(skipped)} skipped", flush=True)
            except Exception as exc:                    # noqa: BLE001
                nfail += 1
                print(f"[{i}/{len(gribs)}] FAIL {os.path.basename(futs[fut])}: {exc}",
                      flush=True)

    print(f"DONE group={args.group} written={nw} skipped={ns} failed_files={nfail}")
    return 1 if nfail else 0


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