#!/usr/bin/env python3
"""
Fetch NOAA GFSWave Arctic 9km GRIB2 and extract wave height, direction, period
without pygrib. Uses xarray+cfgrib (ecCodes backend).
"""
import argparse
import json
import os
import sys
import tempfile
from datetime import datetime
from typing import Dict, Any

import numpy as np
import requests
import xarray as xr


def build_url(date_str: str, run: str, fh: int) -> str:
    base = "https://nomads.ncep.noaa.gov/pub/data/nccf/com/gfs/prod"
    fname = f"gfswave.t{run}z.arctic.9km.f{fh:03d}.grib2"
    return f"{base}/gfs.{date_str}/{run}/wave/gridded/{fname}"


def download(url: str) -> str:
    r = requests.get(url, timeout=300)
    r.raise_for_status()
    tf = tempfile.NamedTemporaryFile(delete=False, suffix=".grib2")
    tf.write(r.content)
    tf.close()
    return tf.name


def open_grib(path: str) -> xr.Dataset:
    # cfgrib uses ecCodes; ensure eccodes is installed on the system
    return xr.open_dataset(path, engine="cfgrib", decode_timedelta=True)


def pick_var(ds: xr.Dataset, cands) -> str | None:
    for c in cands:
        if c in ds.variables:
            return c
    # heuristic fallback
    for v in ds.variables:
        lv = v.lower()
        if ("wave" in lv and "height" in lv) or v in ("HTSGW", "htsgw"):
            return v
    return None


def sample_points(ds: xr.Dataset, n: int = 1000) -> list[Dict[str, Any]]:
    lat_name = "latitude" if "latitude" in ds.variables else "lat"
    lon_name = "longitude" if "longitude" in ds.variables else "lon"
    lats = ds[lat_name].values
    lons = ds[lon_name].values

    wave_var = pick_var(ds, ["swh", "HTSGW", "htsgw"])  # significant wave height
    if wave_var is None:
        raise RuntimeError("No wave height variable found in GRIB")
    wh = ds[wave_var].values

    dir_var = pick_var(ds, [
        "dirpw", "DIRPW", "dp", "wvdir", "WVDIR", "dir", "mwd", "MWD", "MWDIR",
    ])
    per_var = pick_var(ds, [
        "perpw", "PERPW", "tp", "wvper", "WVPER", "per", "pp1d", "PP1D", "mwp", "MWP",
    ])
    wd = ds[dir_var].values if dir_var else None
    wp = ds[per_var].values if per_var else None

    lon_grid, lat_grid = np.meshgrid(lons, lats)
    flat_lats = lat_grid.flatten()
    flat_lons = lon_grid.flatten()
    flat_wh = wh.flatten()

    mask = np.isfinite(flat_wh)
    if wd is not None:
        mask &= np.isfinite(wd.flatten())
    idx_all = np.where(mask)[0]
    if idx_all.size == 0:
        return []
    choose = np.random.choice(idx_all, size=min(n, idx_all.size), replace=False)

    out = []
    for i in choose:
        item: Dict[str, Any] = {
            "lat": float(flat_lats[i]),
            "lon": float(flat_lons[i]),
            "wave_height": float(flat_wh[i]),
        }
        if wd is not None:
            item["wave_direction"] = float(wd.flatten()[i])
        if wp is not None:
            item["wave_period"] = float(wp.flatten()[i])
        out.append(item)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--date", help="YYYYMMDD (default: today UTC)")
    ap.add_argument("--run", default=None, help="Cycle hour: 00,06,12,18 (default: best guess)")
    ap.add_argument("--fh", type=int, default=0, help="Forecast hour (0-240)")
    ap.add_argument("--out-json", default=None, help="Write sampled points JSON")
    # ap.add_argument("--out-nc", default=None, help="Write NetCDF copy of the GRIB")
    args = ap.parse_args()

    now = datetime.utcnow()
    date_str = args.date or now.strftime("%Y%m%d")
    # pick a likely available run by current hour if not provided
    if args.run is None:
        hr = now.hour
        if hr >= 18:
            run = "18"
        elif hr >= 12:
            run = "12"
        elif hr >= 6:
            run = "06"
        else:
            run = "00"
    else:
        run = args.run.zfill(2)

    url = build_url(date_str, run, args.fh)
    print(f"Downloading: {url}")
    path = download(url)
    print(f"Saved: {path}")

    ds = open_grib(path)
    pts = sample_points(ds, n=1000)

    # Summaries
    wh_name = pick_var(ds, ["swh", "HTSGW", "htsgw"]) or "swh"
    wh = ds[wh_name].values
    print(json.dumps({
        "date": date_str,
        "run": run,
        "forecast_hour": args.fh,
        "points": len(pts),
        "height_min": float(np.nanmin(wh)),
        "height_max": float(np.nanmax(wh)),
        "height_mean": float(np.nanmean(wh)),
    }, indent=2))

    if args.out_json:
        with open(args.out_json, "w") as f:
            json.dump({"type": "points", "points": pts}, f)
        print(f"Wrote JSON: {args.out_json}")

    # if args.out_nc:
    #     ds.to_netcdf(args.out_nc)
    #     print(f"Wrote NetCDF: {args.out_nc}")

    ds.close()

    # Clean up temp file
    try:
        os.unlink(path)
    except Exception:
        pass


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