#!/usr/bin/env python3
"""Fetch CAISO day-ahead (DAM) + real-time (RTM) hourly LMP for one hub (free OASIS, no key).

The RT-vs-DA spread is the forecast-error-driven, less-seasonal intraday opportunity. Monthly
chunks with backoff (OASIS rate-limits hard). Aggregates RT to hourly. Caches per chunk.
Saves data/surfer/caiso/{dam,rtm}.json as {hour_key: price}.
"""
import io
import json
import os
import time
import urllib.request
import zipfile

NODE = "TH_NP15_GEN-APND"
OUT = "data/surfer/caiso"
MONTHS = [(2024, m) for m in range(1, 13)]


def oasis(qn, market, y, m):
    s = f"{y}{m:02d}01T08:00-0000"
    ny, nm = (y + 1, 1) if m == 12 else (y, m + 1)
    e = f"{ny}{nm:02d}01T08:00-0000"
    u = (f"http://oasis.caiso.com/oasisapi/SingleZip?queryname={qn}&startdatetime={s}"
         f"&enddatetime={e}&version=1&market_run_id={market}&node={NODE}&resultformat=6")
    for a in range(5):
        try:
            raw = urllib.request.urlopen(urllib.request.Request(u, headers={"User-Agent": "curl/8"}), timeout=90).read()
            z = zipfile.ZipFile(io.BytesIO(raw))
            txt = z.read(z.namelist()[0]).decode()
            return txt
        except Exception as ex:
            time.sleep(12 * (a + 1))
    return None


def parse_hourly(txt):
    lines = txt.strip().split("\n")
    hdr = lines[0].split(",")
    if "LMP_TYPE" not in hdr:
        return {}                                              # error report / wrong format
    iT = hdr.index("LMP_TYPE")
    iV = hdr.index("MW") if "MW" in hdr else (hdr.index("PRC") if "PRC" in hdr else -1)
    if iV < 0:
        return {}
    iD = hdr.index("OPR_DT"); iH = hdr.index("OPR_HR")
    agg = {}
    for ln in lines[1:]:
        c = ln.split(",")
        if len(c) <= max(iT, iV, iD, iH):
            continue
        if c[iT] != "LMP":
            continue
        k = f"{c[iD]}H{int(c[iH]):02d}"
        agg.setdefault(k, []).append(float(c[iV]))
    return {k: sum(v) / len(v) for k, v in agg.items()}


def fetch(market, qn):
    os.makedirs(OUT, exist_ok=True)
    out = {}
    for y, m in MONTHS:
        cf = f"{OUT}/{market}_{y}{m:02d}.json"
        if os.path.exists(cf):
            out.update(json.load(open(cf))); continue
        txt = oasis(qn, market, y, m)
        if not txt:
            print(f"  {market} {y}-{m:02d}: FAILED"); continue
        h = parse_hourly(txt)
        json.dump(h, open(cf, "w")); out.update(h)
        print(f"  {market} {y}-{m:02d}: {len(h)} hrs")
        time.sleep(6)
    return out


def main():
    dam = fetch("DAM", "PRC_LMP"); rtm = fetch("RTPD", "PRC_RTPD_LMP")   # 15-min RT pre-dispatch
    json.dump(dam, open(f"{OUT}/dam.json", "w")); json.dump(rtm, open(f"{OUT}/rtm.json", "w"))
    print(f"DONE: DAM {len(dam)} hrs, RTM {len(rtm)} hrs")


if __name__ == "__main__":
    main()