"""FlowAi API example (Python 3.9+, requests): a gauge's latest reading, its observations and forecast as CSV,
and several gauges in one call through the PI-JSON feed.

    pip install requests
    export FLOWAI_API_KEY="fai_..."        # from Profile > Integrations; never put the key in the script
    python flowai_example.py
"""
import csv
import os
import time
from datetime import datetime, timedelta, timezone

import requests

API = os.environ.get("FLOWAI_API_URL", "https://api.flowai-forecast.com")
GAUGE = "02KF019"                       # a WSC station id
GAUGES = ["02KF019", "02HB029"]

session = requests.Session()
session.headers["Authorization"] = f"Bearer {os.environ['FLOWAI_API_KEY']}"


def get(path, params=None):
    """GET a FlowAi URL. On the key's rate limit (HTTP 429), wait as long as Retry-After says, then try once more."""
    r = session.get(f"{API}{path}", params=params, timeout=30)
    if r.status_code == 429:
        time.sleep(int(r.headers.get("Retry-After", "60")))
        r = session.get(f"{API}{path}", params=params, timeout=30)
    r.raise_for_status()
    return r


# 1. The gauge and its latest reading
gauge = get(f"/api/v1/gauges/{GAUGE}").json()
latest = get(f"/api/v1/gauges/{GAUGE}/latest").json()
print(f"{gauge['name']} ({GAUGE}): {latest['flow_cms']} m3/s at {latest['flow_timestamp']}")

# 2. Daily observed flow and level, last 7 days, newest first, to CSV. Each row is the provider's daily value,
#    stamped with its calendar date (00:00). For real-time readings use /recent (the last 48 hours) or the feed (step 4).
observed = get(f"/api/v1/gauges/{GAUGE}/observed", {"days": 7}).json()["readings"]
with open(f"{GAUGE}_observed.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["date", "flow_cms", "level_m", "flag"])
    for row in observed:
        w.writerow([row["timestamp"][:10], row["flow_cms"], row["level_m"], row["flag"]])
print(f"{len(observed)} days observed -> {GAUGE}_observed.csv")

# 3. The latest forecast, one row per day, to CSV
forecast = get(f"/api/v1/gauges/{GAUGE}/forecast").json()
days = [lead for lead in forecast["forecasts"] if lead["segment"] == "forecast"]
with open(f"{GAUGE}_forecast.csv", "w", newline="") as f:
    w = csv.writer(f)
    w.writerow(["issue_time_utc", "valid_date", "lead_days", "flow_cms"])
    for lead in days:
        w.writerow([forecast["issue_time"], lead["valid_date"], lead["lead_days"], lead["flow_cms"]])
print(f"forecast issued {forecast['issue_time']}: {len(days)} days -> {GAUGE}_forecast.csv")

# 4. Several gauges and series in one call: the PI-JSON feed
now = datetime.now(timezone.utc)
params = [("locationIds", g) for g in GAUGES] + [("parameterIds", p) for p in ("Q.obs", "Q.fcst.day")]
params += [("startTime", (now - timedelta(days=2)).strftime("%Y-%m-%dT%H:%M:%SZ")),
           ("endTime", (now + timedelta(days=11)).strftime("%Y-%m-%dT%H:%M:%SZ")),
           ("documentFormat", "PI_JSON")]
feed = get("/rest/fewspiservice/v1/timeseries", params)
for series in feed.json()["timeSeries"]:
    h, events = series["header"], series["events"]
    last = f"{events[-1]['date']} {events[-1]['time']} = {events[-1]['value']}" if events else "no values"
    print(f"{h['locationId']} {h['parameterId']}: {len(events)} values, last {last}")
print(feed.headers["X-FlowAi-Disclaimer"])