""".. title:: WPC Outlook Service

Return to `API Services </api/>`_.

Documentation for /json/wpcoutlook.py
-------------------------------------

This service provides access to the Weather Prediction Center's Excessive
Rainfall Outlook products.
The service is designed to be called with a latitude and longitude point.

Changelog
---------

- 2026-02-24: For IEM consistency, the `generation_time` root attribute was
  renamed `generated_at`.
- 2025-06-29: Initial implementation of WPC Outlook service.

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

Get all the day 1 outlooks for Pierre, South Dakota.  First in JSON, then
CSV, and finally Excel.

https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=44.368&lon=-100.336&day=1
https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=44.368&lon=-100.336&day=1&fmt=csv
https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=44.368&lon=-100.336&day=1&fmt=excel

Provide the day 1 outlook for Pierre, SD valid at 8 UTC on 3 Aug 2024

https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=44.368&lon=-100.336&day=1&time=2024-08-03T08:00

Same request, but for 8 AM CDT

https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=44.368&lon=-100.336&day=1&time=2024-08-03T08:00-05:00


Get only the last day 2 outlook for Washington, DC

https://mesonet.agron.iastate.edu/json/wpcoutlook.py\
?lat=38.907&lon=-77.037&day=2&last=1

"""

import json
from datetime import datetime, timedelta, timezone
from io import BytesIO
from typing import Annotated

import pandas as pd
from pydantic import Field
from pyiem.database import get_sqlalchemy_conn, sql_helper
from pyiem.exceptions import IncompleteWebRequest
from pyiem.nws.products.spcpts import THRESHOLD_ORDER
from pyiem.reference import ISO8601
from pyiem.util import utc
from pyiem.webutil import CGIModel, iemapp

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


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

    callback: CALLBACK_FIELD = None
    fmt: Annotated[
        str,
        Field(
            description="The format to return data in, json, excel, or csv",
            pattern="^(json|excel|csv)$",
        ),
    ] = "json"
    lat: LATITUDE_FIELD = 42.0
    lon: LONGITUDE_FIELD = -95.0
    last: Annotated[
        int, Field(description="Limit to last N outlooks, 0 for all", ge=0)
    ] = 0
    day: OUTLOOK_DAY_FIELD = 1
    time: Annotated[
        str | None,
        Field(
            description=(
                "Optional specification for a valid timestamp to query "
                "outlooks for.  "
                "This is either a ISO8601 timestamp or 'current' for now."
            ),
        ),
    ] = None


def get_order(threshold):
    """Lookup a threshold and get its rank, higher is more extreme"""
    if threshold not in THRESHOLD_ORDER:
        return -1
    return THRESHOLD_ORDER.index(threshold)


def process_df(outlooks: pd.DataFrame) -> pd.DataFrame:
    """Condition the dataframe."""
    if outlooks.empty:
        return outlooks
    outlooks["threshold_rank"] = outlooks["threshold"].apply(get_order)
    return outlooks


def dotime(time: str, lon, lat, day) -> tuple[pd.DataFrame, datetime]:
    """Query for Outlook based on some timestamp"""
    if time in ["", "current", "now"]:
        ts = utc()
        if day > 1:
            ts += timedelta(days=day - 1)
    else:
        try:
            ts = datetime.fromisoformat(time.replace("Z", "+00:00"))
        except ValueError as exp:
            raise IncompleteWebRequest(
                "time must be a valid ISO8601 timestamp"
            ) from exp
        if ts.tzinfo is None:
            ts = ts.replace(tzinfo=timezone.utc)
        else:
            ts = ts.astimezone(timezone.utc)
    with get_sqlalchemy_conn("postgis") as conn:
        outlooks = pd.read_sql(
            sql_helper("""
        SELECT issue at time zone 'UTC' as i,
        expire at time zone 'UTC' as e,
        product_issue at time zone 'UTC' as v,
        threshold, category from spc_outlooks where
        product_issue = (
            select product_issue from spc_outlook where
            issue <= :ts and expire > :ts and day = :day
            and outlook_type = 'E' ORDER by product_issue DESC LIMIT 1)
        and ST_Contains(geom, ST_Point(:lon, :lat, 4326))
        and day = :day and outlook_type = 'E' and category = 'CATEGORICAL'
        """),
            conn,
            params={"lon": lon, "lat": lat, "day": day, "ts": ts},
            index_col=None,
        )
    return process_df(outlooks), ts


def dowork(lon, lat, day) -> pd.DataFrame:
    """Actually do stuff"""
    with get_sqlalchemy_conn("postgis") as conn:
        # Need to compute SIGN seperately
        outlooks = pd.read_sql(
            sql_helper("""
        WITH data as (
            SELECT issue at time zone 'UTC' as i,
            expire at time zone 'UTC' as e,
            product_issue at time zone 'UTC' as v,
            o.threshold, category, t.priority,
            row_number() OVER (PARTITION by expire
                ORDER by priority DESC NULLS last, issue ASC) as rank
            from spc_outlooks o, spc_outlook_thresholds t
            where o.threshold = t.threshold and
            ST_Contains(geom, ST_Point(:lon, :lat, 4326))
            and day = :day and outlook_type = 'E'
            and o.threshold not in ('TSTM', 'SIGN') ORDER by issue DESC),
        agg as (
            select i, e, v, threshold, category from data where rank = 1),
        sign as (
            SELECT issue at time zone 'UTC' as i,
            expire at time zone 'UTC' as e,
            product_issue at time zone 'UTC' as v,
            threshold, category from spc_outlooks
            where ST_Contains(geom, ST_Point(:lon, :lat, 4326))
            and day = :day and outlook_type = 'E'
            and threshold = 'SIGN' ORDER by expire DESC, issue ASC LIMIT 1)

        (SELECT i, e, v, threshold, category from agg
        ORDER by e DESC, threshold desc) UNION ALL
        (SELECT i, e, v, threshold, category from sign
        ORDER by e DESC, threshold desc)
        """),
            conn,
            params={"lon": lon, "lat": lat, "day": day},
        )
    return outlooks


@iemapp(help=__doc__, schema=Schema)
def application(environ: dict, start_response: callable):
    """Answer request."""
    time = environ.get("time")
    fmt = environ["fmt"]
    lon = environ["lon"]
    lat = environ["lat"]
    last = environ["last"]
    day = environ["day"]
    ts = None
    if time is not None:
        outlooks, ts = dotime(time, lon, lat, day)
        if not outlooks.empty:
            outlooks = outlooks.iloc[[0]]
    else:
        outlooks = dowork(lon, lat, day)

    if fmt == "json" and time is not None:
        res = json_response_dict(
            {
                "query_params": {
                    "time": ts.strftime(ISO8601),
                    "lon": lon,
                    "lat": lat,
                    "day": day,
                },
                "outlook": {},
            }
        )
        if not outlooks.empty:
            row0 = outlooks.iloc[0]
            res["outlook"] = {
                "threshold": row0["threshold"],
                "utc_product_issue": pd.Timestamp(row0["v"]).strftime(ISO8601),
                "utc_issue": pd.Timestamp(row0["i"]).strftime(ISO8601),
                "utc_expire": pd.Timestamp(row0["e"]).strftime(ISO8601),
            }
        payload = json.dumps(res)
        start_response("200 OK", [("Content-type", "application/json")])
        return payload
    if fmt == "json":
        running = {}
        res = {"outlooks": []}
        for _, row in outlooks.iterrows():
            if last > 0:
                running.setdefault(row["threshold"], 0)
                running[row["threshold"]] += 1
                if running[row["threshold"]] > last:
                    continue
            res["outlooks"].append(
                dict(
                    day=day,
                    utc_issue=row["i"].strftime(ISO8601),
                    utc_expire=row["e"].strftime(ISO8601),
                    utc_product_issue=row["v"].strftime(ISO8601),
                    threshold=row["threshold"],
                    category=row["category"],
                )
            )
        payload = json.dumps(res)
        start_response("200 OK", [("Content-type", "application/json")])
        return payload

    outlooks = outlooks.rename(
        columns={
            "i": "utc_issue",
            "e": "utc_expire",
            "v": "utc_product_issue",
        },
    )
    if fmt == "excel":
        headers = [
            ("Content-type", "application/vnd.ms-excel"),
            ("Content-Disposition", "attachment; filename=outlooks.xls"),
        ]
        start_response("200 OK", headers)
        with BytesIO() as bio:
            outlooks.to_excel(bio, index=False)
            return bio.getvalue()

    headers = [
        ("Content-type", "text/csv"),
        ("Content-Disposition", "attachment; filename=outlooks.csv"),
    ]
    payload = outlooks.to_csv(index=False)
    start_response("200 OK", headers)
    return payload
