"""
Use the DWD ICON model to estimate IEMRE hourly fields for non CONUS domains.

https://www.dwd.de/DE/leistungen/opendata/help/modelle/Opendata_cdo_EN.pdf
https://www.dwd.de/DWD/forschung/nwv/fepub/icon_database_main.pdf

Orphaned for now, use_ifs.py instead.

"""

import os
import pathlib
import subprocess
import sys
import tempfile
import time
from datetime import datetime, timedelta, timezone

import click
import httpx
import numpy as np
import pygrib
from affine import Affine
from pyiem.iemre import (
    DOMAINS,
    get_hourly_ncname,
    hourly_offset,
    reproject2iemre,
)
from pyiem.util import logger, ncopen

LOG = logger()
META = {
    "skyc": {"gname": "clct"},
    "tmpk": {"gname": "t_2m"},
    "dwpk": {"gname": "td_2m"},
    "uwnd": {"gname": "u_10m"},
    "vwnd": {"gname": "v_10m"},
    "p01m": {"gname": "tot_prec"},
    "p01m_prev": {"gname": "tot_prec", "offset": 1},
    # Inconsistent on server, punted "soil4t": {"gname": "t_so"},
    # This is net shortwave
    "rsds": {"gname": "asob_s"},
    "rsds_prev": {"gname": "asob_s", "offset": 1},
}


def compute_model_valid(valid: datetime) -> datetime | None:
    """
    Compute the model valid time based on the provided valid datetime.
    The ICON model data is available at 00, 06, 12, and 18 UTC.
    """
    # We have to avoid F000 as precip and solar do not exist
    for offset in range(1, 24):
        model_valid = valid - timedelta(hours=offset)
        if model_valid.hour % 6 != 0:
            continue
        testfn = (
            f"https://opendata.dwd.de/weather/nwp/icon/grib/{model_valid:%H}/"
            f"t_2m/icon_global_icosahedral_single-level_{model_valid:%Y%m%d%H}_"
            f"{offset:03.0f}_T_2M.grib2.bz2"
        )
        try:
            with httpx.Client() as client:
                response = client.head(testfn)
            if response.status_code == 200:
                LOG.info("Found ICON model data for %s", model_valid)
                return model_valid
        except httpx.RequestError:
            # Handle request errors (e.g., network issues)
            continue
    return None


def grib_download(model_valid: datetime, valid: datetime) -> None:
    """
    Download the necessary GRIB files for the ICON model.
    """
    baseurl = "https://opendata.dwd.de/weather/nwp/icon/grib/"
    fhour = (valid - model_valid).total_seconds() // 3600
    for var, meta in META.items():
        filename = f"{var}.grib2.bz2"
        fhour_off = fhour - meta.get("offset", 0)
        url = (
            f"{baseurl}{model_valid:%H}/{meta['gname']}/"
            "icon_global_icosahedral_single-level_"
            f"{model_valid:%Y%m%d%H}_{fhour_off:03.0f}_{meta['gname'].upper()}"
            ".grib2.bz2"
        )
        LOG.info("Downloading %s", url)
        for attempt in range(3):
            try:
                with httpx.Client() as client:
                    response = client.get(url)
                    if response.status_code == 404:
                        # Try something 6 hours older
                        mv2 = model_valid - timedelta(hours=6)
                        fo2 = fhour_off + 6
                        url = (
                            f"{baseurl}{mv2:%H}/{meta['gname']}/"
                            "icon_global_icosahedral_single-level_"
                            f"{mv2:%Y%m%d%H}_{fo2:03.0f}_"
                            f"{meta['gname'].upper()}.grib2.bz2"
                        )
                        response = client.get(url)
                response.raise_for_status()
                break
            except (httpx.RequestError, httpx.HTTPStatusError) as e:
                LOG.error("Failed to download %s: %s", filename, e)
                time.sleep(5)
            if attempt == 2:
                LOG.warning("Aborting")
                sys.exit(3)

        pathlib.Path(filename).write_bytes(response.content)
        subprocess.run(["bzip2", "-d", filename], check=True)
        # Use cdo to convert grid to lat/lon
        with subprocess.Popen(
            [
                "cdo",
                "-f",
                "grb2",
                (
                    "remap,/mesonet/data/meta/icon_description.txt,"
                    "/mesonet/data/meta/icon_weights.nc"
                ),
                f"{var}.grib2",
                f"{var}_latlon.grib2",
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        ) as proc:
            stdout, stderr = proc.communicate()
        if proc.returncode != 0:
            LOG.error(
                "CDO remap failed for %s: %s %s",
                filename,
                stderr.decode("utf-8"),
                stdout.decode("utf-8"),
            )


def copy_grib_to_netcdf(valid: datetime, domain: str, fhour: int) -> None:
    """Fun times."""
    idx = hourly_offset(valid)
    # grib file is stored S to N
    affine_in = Affine(0.125, 0, -180.0, 0, 0.125, -90.0)
    with ncopen(get_hourly_ncname(valid.year, domain), "a") as nc:
        for var in META:
            if var.endswith("_prev"):  # lame
                continue
            grib_file = f"{var}_latlon.grib2"
            if not os.path.exists(grib_file):
                LOG.warning("GRIB file %s does not exist, skipping", grib_file)
                continue
            with pygrib.open(grib_file) as grbs:
                values = grbs[1].values
            if var == "p01m":
                with pygrib.open(f"{var}_prev_latlon.grib2") as grbs:
                    values = values - grbs[1].values
            if var == "rsds":
                with pygrib.open(f"{var}_prev_latlon.grib2") as grbs:
                    values = (values * fhour) - (grbs[1].values * (fhour - 1))
            iemre_data = reproject2iemre(
                values, affine_in, "EPSG:4326", domain=domain
            )
            LOG.info(
                "Write %s[%s]@%s mean: %.2f",
                var,
                domain,
                idx,
                np.mean(iemre_data),
            )
            nc.variables[var][idx] = iemre_data


@click.command()
@click.option("--valid", type=click.DateTime(), required=True)
def main(valid: datetime) -> None:
    """Main function to process ICON data for IEMRE."""
    valid = valid.replace(tzinfo=timezone.utc)
    # 1. Figure out which ICON model is available for usage
    model_valid = compute_model_valid(valid)
    if model_valid is None:
        LOG.warning("No ICON model data available for %s", valid)
        return
    # 2. Download the necessary files
    grib_download(model_valid, valid)
    fhour = (valid - model_valid).total_seconds() // 3600
    # 3. Copy to IEMRE netcdf files
    for domain in DOMAINS:
        if domain == "conus":
            continue
        copy_grib_to_netcdf(valid, domain, fhour)


if __name__ == "__main__":
    with tempfile.TemporaryDirectory() as tmpdir:
        os.chdir(tmpdir)
        main()
