#!/usr/bin/env python3
"""
CDS API helper — fetches ERA5/CMIP6 data and prints JSON to stdout.
Called by weather_cds tool with a JSON request on stdin.

Input (stdin, JSON):
  {
    "url":      "https://cds.climate.copernicus.eu/api",
    "key":      "<personal-access-token>",
    "dataset":  "reanalysis-era5-single-levels",
    "request":  { ... CDS request dict ... },
    "lat":      37.5,
    "lon":      126.9
  }

Output (stdout, JSON):
  { "ok": true, "data": { variable: {...} }, "times": [...], "meta": {...} }
  { "ok": false, "error": "..." }
"""

import sys, json, os, tempfile, zipfile

def main():
    inp = json.load(sys.stdin)
    url     = inp.get("url", "https://cds.climate.copernicus.eu/api")
    key     = inp.get("key", "")
    dataset = inp.get("dataset", "reanalysis-era5-single-levels")
    request = inp.get("request", {})
    lat     = inp.get("lat")
    lon     = inp.get("lon")

    try:
        import cdsapi
    except ImportError:
        print(json.dumps({"ok": False, "error": "cdsapi 미설치. pip install cdsapi --break-system-packages"}))
        return

    try:
        import netCDF4 as nc4
        import numpy as np
    except ImportError:
        print(json.dumps({"ok": False, "error": "netCDF4/numpy 미설치. pip install netCDF4 numpy --break-system-packages"}))
        return

    tmpzip = None
    tmpnc  = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f:
            tmpzip = f.name

        client = cdsapi.Client(url=url, key=key, quiet=True, progress=False)
        client.retrieve(dataset, request, tmpzip)

        # CDS now returns a ZIP containing NetCDF + provenance files
        with tempfile.NamedTemporaryFile(suffix=".nc", delete=False) as f:
            tmpnc = f.name

        if zipfile.is_zipfile(tmpzip):
            with zipfile.ZipFile(tmpzip) as z:
                nc_names = [n for n in z.namelist() if n.endswith(".nc")]
                if not nc_names:
                    raise RuntimeError("ZIP에 .nc 파일이 없습니다: " + str(z.namelist()))
                with open(tmpnc, "wb") as out:
                    out.write(z.read(nc_names[0]))
        else:
            # Not a ZIP — assume raw NetCDF
            import shutil
            shutil.copy(tmpzip, tmpnc)

        ds = nc4.Dataset(tmpnc)
        result = {"ok": True, "data": {}, "times": [], "meta": {}}

        # --- time axis ---
        if "time" in ds.variables:
            t_var = ds["time"]
            t_units = getattr(t_var, "units", "")
            t_cal   = getattr(t_var, "calendar", "standard")
            try:
                import cftime
                dts = nc4.num2date(t_var[:], t_units, t_cal)
                result["times"] = [str(d)[:16] for d in dts]
            except Exception:
                result["times"] = [str(v) for v in t_var[:].tolist()]

        # --- data variables (exclude coordinate/bound variables) ---
        COORD_VARS = {"time", "time_bnds", "lat", "lat_bnds", "lon", "lon_bnds",
                      "latitude", "longitude", "height", "level", "plev", "bnds"}
        data_vars = [v for v in ds.variables if v not in COORD_VARS]

        for vname in data_vars:
            var = ds[vname]
            units     = getattr(var, "units", "")
            long_name = getattr(var, "long_name", vname)

            raw = var[:].filled(np.nan) if hasattr(var[:], "filled") else np.array(var[:])

            # Nearest-point extraction when lat/lon given
            if lat is not None and lon is not None:
                lat_dim = lon_dim = None
                for dim in var.dimensions:
                    if dim in ("lat", "latitude"):  lat_dim = dim
                    if dim in ("lon", "longitude"):  lon_dim = dim
                if lat_dim and lon_dim:
                    lat_idx = int(np.argmin(np.abs(ds[lat_dim][:] - lat)))
                    lon_idx = int(np.argmin(np.abs(ds[lon_dim][:] - lon)))
                    slices = tuple(
                        lat_idx if d == lat_dim else (lon_idx if d == lon_dim else slice(None))
                        for d in var.dimensions
                    )
                    raw = raw[slices]

            vals = raw.flatten().tolist()

            # Unit conversions
            if units == "K":
                vals = [round(v - 273.15, 3) if v is not None and not (isinstance(v, float) and v != v) else None for v in vals]
                units = "°C"
            elif units == "m s**-1":
                units = "m/s"
            elif units == "m" and ("precip" in vname.lower() or "pr" == vname):
                vals = [round(v * 1000, 3) if v is not None else None for v in vals]
                units = "mm"

            arr = [v for v in vals if v is not None and isinstance(v, float) and v == v]
            stats = {}
            if arr:
                stats = {"mean": round(float(np.mean(arr)), 3),
                         "min":  round(float(np.min(arr)),  3),
                         "max":  round(float(np.max(arr)),  3),
                         "n":    len(arr)}
            result["data"][vname] = {
                "values": vals, "units": units, "long_name": long_name, "stats": stats
            }

        # lat/lon info
        for dim in ("lat", "latitude"):
            if dim in ds.variables:
                result["meta"]["lat_grid"] = [round(float(v), 4) for v in ds[dim][:].tolist()]
                break
        for dim in ("lon", "longitude"):
            if dim in ds.variables:
                result["meta"]["lon_grid"] = [round(float(v), 4) for v in ds[dim][:].tolist()]
                break

        result["meta"].update({
            "dataset": dataset,
            "variables": data_vars,
            "n_times": len(result["times"]),
        })

        ds.close()
        print(json.dumps(result))

    except Exception as e:
        print(json.dumps({"ok": False, "error": str(e)}))
    finally:
        for f in [tmpzip, tmpnc]:
            if f and os.path.exists(f):
                try: os.unlink(f)
                except Exception: pass

if __name__ == "__main__":
    main()