""".. title:: JSON Service for Stage IV Precipitation Data

Return to `API Services </api/#json>`_

Documentation for /json/stage4.py
---------------------------------

This emits hourly precipitation data from the Stage IV precipitation dataset
for a given date, which defaults to UTC, but can be specified by the given
`tz` parameter.  Please note that this is not a pure stage IV service, but
includes some bias correcting that the IEM does against PRISM.

Changelog
---------

- 2024-08-19: Added support for tz parameter

Example Usage
-------------

Provide hourly stage IV estimates for 13 August 2024 for 102.3W, 45.1N for
a date in US/Central timezone.

https://mesonet.agron.iastate.edu/json/stage4.py\
?lon=-102.3&lat=45.1&valid=2024-08-13&tz=America/Chicago

"""

import json
import os
from datetime import date, datetime, timedelta
from typing import Annotated
from zoneinfo import ZoneInfo

import numpy as np
from pydantic import Field
from pyiem.grid.nav import get_nav
from pyiem.iemre import hourly_offset
from pyiem.util import mm2inch, ncopen
from pyiem.webutil import CGIModel, iemapp

from iemweb.fields import (
    CALLBACK_FIELD,
    LATITUDE_FIELD,
    LONGITUDE_FIELD,
    TZ_FIELD,
)
from iemweb.util import json_response_dict


class Schema(CGIModel):
    """See how we are called."""

    callback: CALLBACK_FIELD = None
    lat: LATITUDE_FIELD
    lon: LONGITUDE_FIELD
    valid: Annotated[date, Field(description="Valid date of data")]
    tz: TZ_FIELD = "UTC"


def myrounder(val, precision):
    """round a float or give back None"""
    if val is None or np.isnan(val) or np.ma.is_masked(val):
        return None
    return round(val, precision)


def dowork(query: Schema):
    """Do work!"""
    valid = query.valid
    # We want data for the UTC date and timestamps are in the rears, so from
    # 1z through 1z
    sts = datetime(
        valid.year, valid.month, valid.day, 1, tzinfo=ZoneInfo(query.tz)
    )
    ets = sts + timedelta(hours=24)
    sidx = hourly_offset(sts)
    eidx = hourly_offset(ets)

    ncfn = f"/mesonet/data/stage4/{valid:%Y}_stage4_hourly.nc"
    res = json_response_dict(
        {
            "gridi": -1,
            "gridj": -1,
            "for_date_in_timezone": query.tz,
            "data": [],
        }
    )
    if not os.path.isfile(ncfn):
        return json.dumps(res)
    i, j = get_nav("STAGE4").find_ij(query.lon, query.lat)
    if i is not None:
        with ncopen(ncfn) as nc:
            res["gridi"] = i
            res["gridj"] = j

            ppt = nc.variables["p01m"][sidx:eidx, j, i]

        for tx, pt in enumerate(ppt):
            valid = sts + timedelta(hours=tx)
            utcnow = valid.astimezone(ZoneInfo("UTC"))
            res["data"].append(
                {
                    "end_valid": utcnow.strftime("%Y-%m-%dT%H:00:00Z"),
                    "precip_in": myrounder(mm2inch(pt), 2),
                }
            )

    return json.dumps(res)


def get_mckey(environ: dict):
    """Get the memcachekey."""
    return (
        f"/json/stage4/{environ['lon']:.2f}/{environ['lat']:.2f}/"
        f"{environ['valid']}/{environ['tz']}"
    )


@iemapp(
    help=__doc__, schema=Schema, memcachekey=get_mckey, memcacheexpire=3600
)
def application(environ, start_response):
    """Answer request."""
    res = dowork(environ["_cgimodel_schema"])
    start_response("200 OK", [("Content-type", "application/json")])
    return res
