#!/usr/bin/env python3
"""Fetch full-history ES front-month OHLCV-1m (continuous .c.0), year-chunked, BUDGET-CAPPED.

get_cost-gated (aborts over cap, no download); year chunks with quarter fallback on 504.
Saves per-chunk DBN to data/surfer/es1m/ (gitignored). Key from env, never printed.
"""
import os
import sys
import time

import databento as db

CAP_USD = 25.00
DS = "GLBX.MDP3"
SCHEMA = "ohlcv-1m"
SYM = "ES.c.0"
STYPE = "continuous"
OUT_DIR = "data/surfer/es1m"


def fetch(client, start, end, out):
    data = client.timeseries.get_range(dataset=DS, symbols=[SYM], schema=SCHEMA,
                                       start=start, end=end, stype_in=STYPE)
    data.to_file(out)
    return sum(1 for _ in data)


def main():
    key = os.environ.get("DATABENTO_API_KEY")
    if not key:
        print("DATABENTO_API_KEY not set"); return 2
    client = db.Historical(key)
    cost = client.metadata.get_cost(dataset=DS, symbols=[SYM], schema=SCHEMA,
                                    start="2010-06-06", end="2026-06-05", stype_in=STYPE)
    print(f"aggregate get_cost=${cost:.4f}  cap=${CAP_USD:.2f}")
    if cost > CAP_USD:
        print("ABORT: over cap — nothing downloaded."); return 1
    os.makedirs(OUT_DIR, exist_ok=True)
    total = 0
    for year in range(2010, 2027):
        y0, y1 = f"{year}-01-01", f"{year+1}-01-01"
        if year == 2010:
            y0 = "2010-06-06"
        if year == 2026:
            y1 = "2026-06-05"
        out = f"{OUT_DIR}/ES_{year}.dbn"
        if os.path.exists(out):
            print(f"  {year}: exists, skip"); continue
        ok = False
        for attempt in range(1, 3):
            try:
                n = fetch(client, y0, y1, out); total += n
                print(f"  {year}: {n:,} recs -> {out}")
                ok = True; break
            except Exception as e:
                print(f"  {year}: attempt {attempt} {type(e).__name__}; retry")
                time.sleep(3)
        if not ok:  # fall back to quarter chunks
            print(f"  {year}: year failed → quarter fallback")
            for q, (m0, m1) in enumerate([("01-01", "04-01"), ("04-01", "07-01"),
                                          ("07-01", "10-01"), ("10-01", "12-31")], 1):
                qs, qe = f"{year}-{m0}", f"{year}-{m1}"
                if year == 2010 and q == 1:
                    qs = "2010-06-06"
                if year == 2026 and q >= 3:
                    continue
                qout = f"{OUT_DIR}/ES_{year}_q{q}.dbn"
                if os.path.exists(qout):
                    continue
                try:
                    n = fetch(client, qs, qe, qout); total += n
                    print(f"    {year} q{q}: {n:,} -> {qout}")
                except Exception as e:
                    print(f"    {year} q{q}: FAILED {type(e).__name__}")
                time.sleep(1)
    print(f"DONE: {total:,} records into {OUT_DIR}/")
    return 0


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