#!/usr/bin/env python3
"""Fetch operational CFS TMAX from NOMADS for one initial time, by byte range.

Only the TMAX record is wanted, and every GRIB2 file on NOMADS ships a .idx
sidecar giving each record's byte offset. So the whole 46-day forecast costs
~6 MB instead of ~15 GB: fetch the small .idx, find the TMAX line, then ask for
exactly that byte range. Offsets are data-dependent, so the .idx of EVERY file
has to be read -- an offset taken from one file does not apply to the next.

NOMADS answers HTTP/2 with a PROTOCOL_ERROR on these paths; the requests here go
over HTTP/1.1, which works.

Output: <outdir>/<VALIDTIME>.grb2, one single-record file per 6-hourly step.
"""

from __future__ import annotations

import argparse
import concurrent.futures as cf
import datetime as dt
import os
import sys
import time
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C

BASE = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/cfs/prod"
WANT = ":TMAX:2 m above ground:"
UA = {"User-Agent": "curl/7.87", "Connection": "close"}


def _get(url, rng=None, timeout=120, tries=4):
    for k in range(tries):
        try:
            req = urllib.request.Request(url, headers=dict(UA))
            if rng:
                req.add_header("Range", "bytes=%d-%d" % rng)
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as exc:
            if k == tries - 1:
                raise
            time.sleep(2 * (k + 1))
    return None


def one_step(args):
    init, member, valid, outdir = args
    out = os.path.join(outdir, "%s.grb2" % valid)
    if os.path.exists(out) and os.path.getsize(out) > 1000:
        return valid, "skip"
    d = "cfs.%s/%s/6hrly_grib_%s" % (init[:8], init[8:10], member)
    fn = "flxf%s.%s.%s.grb2" % (valid, member, init)
    url = "%s/%s/%s" % (BASE, d, fn)
    try:
        inv = _get(url + ".idx").decode().splitlines()
    except Exception as exc:
        return valid, "idx loi: %s" % exc
    off = None
    for i, ln in enumerate(inv):
        if WANT in ln:
            a = int(ln.split(":")[1])
            b = (int(inv[i + 1].split(":")[1]) - 1) if i + 1 < len(inv) else None
            off = (a, b if b is not None else a + 4_000_000)
            break
    if off is None:
        return valid, "khong thay TMAX trong idx"
    try:
        blob = _get(url, rng=off)
    except Exception as exc:
        return valid, "tai loi: %s" % exc
    if not blob or blob[:4] != b"GRIB":
        return valid, "khong phai GRIB"
    tmp = out + ".tmp"
    open(tmp, "wb").write(blob)
    os.replace(tmp, out)
    return valid, "ok"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", default="2026082900")
    ap.add_argument("--member", default="01")
    ap.add_argument("--outdir", default=os.path.join(C.OUT, "RT", "cfs_2026082900"))
    ap.add_argument("--nthread", type=int, default=8)
    a = ap.parse_args()
    os.makedirs(a.outdir, exist_ok=True)

    d0 = dt.datetime(int(a.init[:4]), int(a.init[4:6]), int(a.init[6:8]),
                     int(a.init[8:10]))
    vts = [(d0 + dt.timedelta(hours=6 * (k + 1))).strftime("%Y%m%d%H")
           for k in range(C.NSTEP)]
    jobs = [(a.init, a.member, v, a.outdir) for v in vts]

    t0, nok, nskip, bad = time.time(), 0, 0, []
    with cf.ThreadPoolExecutor(a.nthread) as ex:
        for n, (v, st) in enumerate(ex.map(one_step, jobs), 1):
            if st == "ok":
                nok += 1
            elif st == "skip":
                nskip += 1
            else:
                bad.append((v, st))
            if n % 40 == 0:
                print("  %d/%d  tai %d  bo qua %d  loi %d  %.0fs"
                      % (n, len(jobs), nok, nskip, len(bad), time.time() - t0),
                      flush=True)
    mb = sum(os.path.getsize(os.path.join(a.outdir, f))
             for f in os.listdir(a.outdir) if f.endswith(".grb2")) / 1e6
    print("xong: %d buoc, %d loi, %.1f MB, %.0f s"
          % (nok + nskip, len(bad), mb, time.time() - t0))
    for v, s in bad[:10]:
        print("  LOI %s: %s" % (v, s))
    return 1 if bad else 0


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