#!/usr/bin/env python3
"""
Download Theta 15:45 option chunks for multiple roots, matching the existing SPY/VIX/VXX window.

Defaults:
- roots: TLT,GLD,IWM,QQQ,USO
- quote-only: skip greeks (then backfill IV/delta via BS estimator)
- window: inferred intersection of existing SPY/VIX/VXX chunk files under RenTech/data/theta_chunks

Example:
  python RenTech/data_pipeline/download_theta_multi_roots.py --workers 3
  python RenTech/data_pipeline/download_theta_multi_roots.py --start-date 2016-04-01 --end-date 2026-04-30
"""

from __future__ import annotations

import argparse
import subprocess
import sys
from datetime import date
from pathlib import Path

import build_theta_dataset as btd

REPO_ROOT = Path(__file__).resolve().parents[2]
THETA_DIR = REPO_ROOT / "RenTech" / "data" / "theta_chunks"
DEFAULT_ROOTS = ("TLT", "GLD", "IWM", "QQQ", "USO")
MATCH_ROOTS = ("SPY", "VIX", "VXX")


def _parse_month_from_name(path: Path) -> tuple[int, int] | None:
    # expected: <root>_1545_YYYY_MM(.parquet|_ivfilled.parquet)
    parts = path.stem.split("_")
    if len(parts) < 4:
        return None
    try:
        y = int(parts[2])
        m = int(parts[3])
        if 1 <= m <= 12:
            return (y, m)
    except ValueError:
        return None
    return None


def _month_to_first_day(ym: tuple[int, int]) -> date:
    return date(ym[0], ym[1], 1)


def _month_to_last_day(ym: tuple[int, int]) -> date:
    y, m = ym
    if m == 12:
        nxt = date(y + 1, 1, 1)
    else:
        nxt = date(y, m + 1, 1)
    from datetime import timedelta

    return nxt - timedelta(days=1)


def infer_existing_window(theta_dir: Path, roots: tuple[str, ...]) -> tuple[date, date]:
    mins: list[date] = []
    maxs: list[date] = []
    for r in roots:
        pats = [f"{r.lower()}_1545_*.parquet", f"{r.upper()}_1545_*.parquet"]
        files: list[Path] = []
        for p in pats:
            files.extend(theta_dir.glob(p))
        yms = [_parse_month_from_name(f) for f in files]
        yms = [x for x in yms if x is not None]
        if not yms:
            raise FileNotFoundError(f"No chunk months found for root={r} in {theta_dir}")
        mn = min(yms)
        mx = max(yms)
        mins.append(_month_to_first_day(mn))
        maxs.append(_month_to_last_day(mx))
    # intersection window across roots
    start = max(mins)
    end = min(maxs)
    if start > end:
        raise ValueError(f"No overlapping month window across roots {roots}: {start} > {end}")
    return start, end


def main() -> None:
    ap = argparse.ArgumentParser(description="Batch Theta chunk download for multiple roots (skip greeks).")
    ap.add_argument("--roots", type=str, default=",".join(DEFAULT_ROOTS), help="Comma list, e.g. TLT,GLD,IWM,QQQ,USO")
    ap.add_argument("--workers", type=int, default=2, help="Parallel month workers per root")
    ap.add_argument("--timeout-sec", type=float, default=600.0)
    ap.add_argument("--start-date", type=str, default="", help="YYYY-MM-DD (optional)")
    ap.add_argument("--end-date", type=str, default="", help="YYYY-MM-DD (optional)")
    ap.add_argument(
        "--fill-greeks-estimator",
        action="store_true",
        help="After download, backfill implied_vol/delta using BS estimator (inplace).",
    )
    ap.add_argument(
        "--match-existing-window",
        action="store_true",
        help="Infer start/end from existing SPY/VIX/VXX chunks (default if no explicit dates provided).",
    )
    args = ap.parse_args()

    roots = tuple(r.strip().upper() for r in args.roots.split(",") if r.strip())
    if not roots:
        raise SystemExit("No roots provided")

    if args.start_date and args.end_date:
        sd = date.fromisoformat(args.start_date)
        ed = date.fromisoformat(args.end_date)
    else:
        sd, ed = infer_existing_window(THETA_DIR, MATCH_ROOTS)

    print(f"Window: {sd} -> {ed}")
    print(f"Roots: {', '.join(roots)}")
    print(f"skip_greeks=True allow_quote_only=True workers={max(1, int(args.workers))} timeout={float(args.timeout_sec)}")

    for root in roots:
        print(f"\n=== Downloading {root} ===", flush=True)
        btd.build_dataset(
            root=root,
            start_date=sd,
            end_date=ed,
            workers=max(1, int(args.workers)),
            allow_quote_only=True,
            skip_greeks=True,
            request_timeout_sec=float(args.timeout_sec),
        )

    if bool(args.fill_greeks_estimator):
        print("\n=== Backfilling IV/Greeks via estimator ===", flush=True)
        cmd = [
            sys.executable,
            str(REPO_ROOT / "RenTech" / "data_pipeline" / "enrich_theta_parquet_iv.py"),
            "--theta-dir",
            str(THETA_DIR),
            "--roots",
            ",".join(roots),
            "--snapshots",
            "1545",
            "--inplace",
            "--yes",
        ]
        print("Running:", " ".join(cmd), flush=True)
        subprocess.run(cmd, check=True)


if __name__ == "__main__":
    main()