#!/usr/bin/env python3
"""Expand the daily futures universe (16y ohlcv-1d, parent) for a diversified multi-asset test.

Budget-capped, get_cost-gated. Saves to data/surfer/<root>.dbn alongside the existing 22 roots.
"""
import os
import sys

import databento as db

CAP = 20.0
DS = "GLBX.MDP3"
START, END = "2010-06-06", "2026-06-05"
ADD = ["6S", "6N", "6M", "HO", "PL", "PA", "UB", "ZL", "ZM", "ZO", "KE", "ZR", "LE", "HE", "GF", "EMD", "NKD"]


def main():
    c = db.Historical(os.environ["DATABENTO_API_KEY"])
    todo = [r for r in ADD if not os.path.exists(f"data/surfer/{r}.dbn")]
    cost = sum(c.metadata.get_cost(dataset=DS, symbols=[r + ".FUT"], schema="ohlcv-1d",
                                   start=START, end=END, stype_in="parent") for r in todo)
    print(f"to fetch: {todo}\naggregate get_cost=${cost:.2f}  cap=${CAP:.2f}")
    if cost > CAP:
        print("ABORT over cap"); return 1
    for r in todo:
        data = c.timeseries.get_range(dataset=DS, symbols=[r + ".FUT"], schema="ohlcv-1d",
                                      start=START, end=END, stype_in="parent")
        data.to_file(f"data/surfer/{r}.dbn")
        print(f"  {r}: saved")
    print(f"DONE: {len(todo)} roots")
    return 0


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