#!/usr/bin/env python3
"""Download daily OHLCV for the surfer universe — continuous front-month, BUDGET-CAPPED.

Re-confirms get_cost (free, no download) and ABORTS if over the cap before spending.
Saves raw DBN to data/surfer/ (gitignored). Reads DATABENTO_API_KEY from env (never printed).
"""
import datetime
import os
import sys

import databento as db

CAP_USD = 40.00          # hard ceiling on cumulative get_cost; covered by $125 free credits ($0 cash)
PER_ROOT_CAP = 10.00     # per-root sanity cap
DS = "GLBX.MDP3"
SCHEMA = "ohlcv-1d"
START = "2010-06-06"
END = datetime.date.today().isoformat()   # fetch through today (dynamic) — was hardcoded; needed for daily forward-track
ROOTS = ["ES", "NQ", "YM", "RTY", "ZN", "ZB", "ZF", "ZT", "6E", "6J", "6B", "6A",
         "6C", "GC", "SI", "HG", "CL", "NG", "RB", "ZC", "ZS", "ZW"]
# parent symbology = all outright expiries per root (fast; continuous chain resolution 504s).
# We build the continuous series ourselves (volume-roll + ratio-adjust) from these.
STYPE = "parent"
OUT_DIR = "data/surfer"


def main():
    key = os.environ.get("DATABENTO_API_KEY")
    if not key:
        print("DATABENTO_API_KEY not set"); return 2
    client = db.Historical(key)

    global END
    try:                                                # clamp END to the dataset's available end (GLBX lags ~1-3d)
        rng = client.metadata.get_dataset_range(dataset=DS)
        avail = (rng.get("end") or "")[:10]
        if avail:
            END = min(END, avail)
    except Exception:
        pass

    os.makedirs(OUT_DIR, exist_ok=True)
    print(f"fetching {SCHEMA} per-root ({len(ROOTS)} roots, {STYPE}) {START}..{END}; "
          f"per-root get_cost gate (≤${PER_ROOT_CAP}), cumulative cap ${CAP_USD}")
    import time
    total_recs = 0
    spent = 0.0
    failed = []
    for r in ROOTS:
        sym = r + ".FUT"
        out_r = f"{OUT_DIR}/{r}.dbn"
        # per-root cost gate (free, no download)
        try:
            c_r = client.metadata.get_cost(dataset=DS, symbols=[sym], schema=SCHEMA,
                                           start=START, end=END, stype_in=STYPE)
        except Exception as e:
            print(f"  {r}: get_cost failed ({type(e).__name__}); skipping"); failed.append(r); continue
        if c_r > PER_ROOT_CAP or spent + c_r > CAP_USD:
            print(f"  {r}: cost ${c_r:.4f} would breach cap (spent ${spent:.2f}) — SKIP"); failed.append(r); continue
        ok = False
        for attempt in range(1, 4):
            try:
                data = client.timeseries.get_range(dataset=DS, symbols=[sym], schema=SCHEMA,
                                                   start=START, end=END, stype_in=STYPE)
                data.to_file(out_r)
                n = sum(1 for _ in data)
                total_recs += n; spent += c_r
                print(f"  {r}: ${c_r:.4f} {n:,} recs -> {out_r}  (cum ${spent:.2f})")
                ok = True
                break
            except Exception as e:
                print(f"  {r}: attempt {attempt} {type(e).__name__}; retry in 3s")
                time.sleep(3)
        if not ok:
            failed.append(r)
    print(f"DONE: {total_recs:,} records, get_cost-cum ${spent:.2f}, "
          f"{len(ROOTS)-len(failed)}/{len(ROOTS)} roots ok"
          + (f"; FAILED/SKIPPED: {failed}" if failed else ""))
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())