"""IMERG Final Run daily (GPM_3IMERGDF V07, NASA GES DISC), 0.1 deg native,
already a clean UTC-calendar-day precipitation total (mm/day) -- no
de-accumulation needed (unlike ERA5-Land). Coverage starts 2000-06-01.

Requires ~/.netrc with a 'urs.earthdata.nasa.gov' entry AND the GES DISC
application approved in the user's Earthdata profile (one-time step at
https://urs.earthdata.nasa.gov/profile/app_authorizations).

Crops to the Vietnam domain immediately and discards the global (3600x1800)
grid to save space.

Output: raw_aux_data/imerg/<year>/imerg_vn_<year>.nc

Usage: python download_imerg.py <start_year> <end_year>
"""
import os
import sys
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, timedelta

import numpy as np
import requests
import xarray as xr

# HDF5 (backing netCDF4/xarray) is NOT thread-safe in these wheel builds --
# concurrent xr.open_dataset() calls from worker threads segfaulted in
# practice. Network download stays parallel; only the HDF5 read is
# serialized behind this lock.
_HDF5_LOCK = threading.Lock()

BASE = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
OUT_DIR = os.path.join(BASE, "raw_aux_data", "imerg")
URL_TMPL = ("https://gpm1.gesdisc.eosdis.nasa.gov/data/GPM_L3/GPM_3IMERGDF.07/"
            "{y}/{m:02d}/3B-DAY.MS.MRG.3IMERG.{y}{m:02d}{d:02d}-S000000-E235959.V07B.nc4")
LON_MIN, LON_MAX, LAT_MIN, LAT_MAX = 100.5, 111.5, 7.5, 24.5
IMERG_START = date(2000, 6, 1)


def fetch_day(d, session, max_retries=5):
    url = URL_TMPL.format(y=d.year, m=d.month, d=d.day)
    for attempt in range(max_retries):
        try:
            r = session.get(url, timeout=25, allow_redirects=True)
            if r.status_code == 404:
                return None
            if r.status_code == 503:
                if attempt == max_retries - 1:
                    print(f"WARN: failed {d}: 503 after {max_retries} tries")
                    return None
                time.sleep(6 * (attempt + 1))
                continue
            r.raise_for_status()
            if r.content[:4] not in (b"\x89HDF", b"\x89HDF"[:4]):
                if b"<html" in r.content[:200].lower():
                    raise RuntimeError("got HTML instead of data -- auth/app-approval issue")
            import tempfile
            with tempfile.NamedTemporaryFile(suffix=".nc4", delete=False) as tf:
                tf.write(r.content)
                tmp_path = tf.name
            with _HDF5_LOCK:
                ds = xr.open_dataset(tmp_path)
                precip = ds["precipitation"].isel(time=0)  # dims (lon, lat)
                sub = precip.sel(lon=slice(LON_MIN, LON_MAX), lat=slice(LAT_MIN, LAT_MAX))
                arr = sub.values.T.astype(np.float32)  # -> (lat, lon)
                lon = sub["lon"].values
                lat = sub["lat"].values
                ds.close()
            os.remove(tmp_path)
            return lon, lat, arr
        except Exception as e:
            if attempt == max_retries - 1:
                print(f"WARN: failed {d}: {e}")
                return None
            time.sleep(3 * (attempt + 1))


def download_year(year):
    out_path = os.path.join(OUT_DIR, str(year), f"imerg_vn_{year}.nc")
    if os.path.exists(out_path):
        print(f"SKIP: {year} already cached")
        return
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    d0 = max(date(year, 1, 1), IMERG_START)
    d1 = date(year, 12, 31)
    if d0 > d1:
        print(f"SKIP: {year} before IMERG start")
        return
    days = [d0 + timedelta(n) for n in range((d1 - d0).days + 1)]
    results = [None] * len(days)
    lon = lat = None
    n_missing = 0

    def worker(idx_day):
        idx, d = idx_day
        s = requests.Session()
        return idx, fetch_day(d, s)

    with ThreadPoolExecutor(max_workers=8) as ex:
        futures = [ex.submit(worker, (i, d)) for i, d in enumerate(days)]
        n_done = 0
        for fut in as_completed(futures):
            idx, res = fut.result()
            results[idx] = res
            n_done += 1
            if res is None:
                n_missing += 1
            elif lon is None:
                lon, lat, _ = res
            if n_done % 60 == 0:
                print(f"PROGRESS: {year} {n_done}/{len(days)} days fetched")
    if lon is None:
        print(f"WARN: {year} entirely missing/failed, skipping write")
        return
    shape2d = None
    for res in results:
        if res is not None:
            shape2d = res[2].shape
            break
    data = np.full((len(days),) + shape2d, np.nan, dtype=np.float32)
    for i, res in enumerate(results):
        if res is not None:
            data[i] = res[2]
    ds = xr.Dataset(
        {"precip": (("time", "lat", "lon"), data)},
        coords={"time": [np.datetime64(d) for d in days], "lat": lat, "lon": lon},
    )
    ds["precip"].attrs = {"units": "mm/day", "source": "IMERG Final Run V07 (GPM_3IMERGDF), UTC calendar day"}
    ds.to_netcdf(out_path, encoding={"precip": {"zlib": True, "complevel": 4}})
    print(f"DONE: {year} written to {out_path}, missing_days={n_missing}/{len(days)}")


if __name__ == "__main__":
    y0, y1 = int(sys.argv[1]), int(sys.argv[2])
    for y in range(y0, y1 + 1):
        download_year(y)