#!/usr/bin/env python3
"""
Build a 3:45 PM ET SPY options dataset from a local ThetaData Terminal (v3).

Why 3:45 PM?
------------
Using 15:45:00 ET avoids the 4:00 PM quote-widening behavior that can distort
execution assumptions in options backtests.

Transport
---------
This script talks **only** to the local Terminal’s **HTTP REST API** (``requests``),
not the ``thetadata`` Python package. Authentication and entitlements are handled
by the Terminal process you log into in the GUI.

Endpoints (see ThetaData HTTP docs)
----------------------------------
* **NBBO at time:** ``GET /v2/bulk_at_time/option/quote``
  ``root=SPY&exp=0&start_date=YYYYMMDD&end_date=YYYYMMDD&ivl=<ms_since_midnight_ET>``
  ``exp=0`` returns all expirations for that root; **must be requested day-by-day**.

* **Greeks / implied vol (for IV + delta):** ``GET /v2/bulk_hist/option/greeks``
  with ``exp=<YYYYMMDD>`` per expiration (``exp=0`` is not supported for greeks
  the same way as quotes), a narrow ``start_time`` / ``end_time`` window around
  15:45 ET, merged by contract key with the quote pull.

Output layout
-------------
Writes one parquet file per month:

    RenTech/data/theta_chunks/{root}_1545_YYYY_MM.parquet

Default ``root`` is **SPY**; use ``--root VXX`` / ``--root VIX`` for vol products (see
``download_theta_vxvix_options_quotes.py`` for a 10-year quote-only batch).

Each file follows the backtester schema:

    [quote_datetime, expiration, strike, right, bid, ask, implied_vol, delta]

Dependencies
------------
``requests``, ``pandas``, ``tqdm``, ``pyarrow`` (for parquet).

Example::

    pip3 install requests pandas tqdm pyarrow
    python RenTech/data_pipeline/build_theta_dataset.py --workers 4
    python RenTech/data_pipeline/build_theta_dataset.py --allow-quote-only
    python RenTech/data_pipeline/build_theta_dataset.py --skip-greeks
    python RenTech/data_pipeline/build_theta_dataset.py --timeout-sec 600
"""

from __future__ import annotations

import argparse
import gc
import math
import time
from concurrent.futures import ProcessPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import date, timedelta
from pathlib import Path
from typing import Any, Iterable

import pandas as pd
import requests
from tqdm import tqdm

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

DEFAULT_ROOT = "SPY"
START_DATE = date(2016, 1, 1)
END_DATE = date.today()

# 15:45:00 ET in milliseconds since midnight.
AT_TIME_MS = (15 * 60 * 60 + 45 * 60) * 1000

# Local Theta Terminal REST base (default install).
THETA_BASE_URL = "http://127.0.0.1:25510"

REPO_ROOT = Path(__file__).resolve().parents[2]
OUT_DIR = REPO_ROOT / "RenTech" / "data" / "theta_chunks"
OUT_DIR.mkdir(parents=True, exist_ok=True)
LOG_FILE = OUT_DIR / "theta_pull_debug.log"

# HTTP
JSON_HEADERS = {"Accept": "application/json"}
# Large greeks responses can exceed 120s; override with --timeout-sec.
REQUEST_TIMEOUT_SEC = 300.0
MAX_RETRIES = 4
RETRY_SLEEP_SEC = 2.5

# Greeks pull: 1-minute buckets around 15:45.
# Theta's bulk greeks endpoint is designed around 1-minute intervals.
GREEKS_IVL_MS = 60_000
# Theta requires: ivl < (end_time - start_time). Use a 2-minute window so we can
# still pick the tick nearest 15:45:00 while keeping ivl=60s.
GREEKS_WINDOW_START_MS = AT_TIME_MS - 60_000
GREEKS_WINDOW_END_MS = AT_TIME_MS + 119_999
# Pause between per-expiration greek calls (many expiries per day).
GREEKS_INTER_EXP_SLEEP_SEC = 0.03

# Set True in each child process after first log (avoids spamming worker_start).
_WORKER_GREEKS_WINDOW_LOGGED = False


def _parquet_num_rows(path: Path) -> int:
    """
    Fast row count using parquet metadata.

    Returns -1 if we can't read metadata for some reason.
    """
    try:
        import pyarrow.parquet as pq

        return int(pq.ParquetFile(path).metadata.num_rows)
    except Exception:
        return -1


# ---------------------------------------------------------------------------
# Helpers — dates & HTTP
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class MonthWindow:
    start: date
    end: date


def _yyyymmdd(d: date) -> int:
    return d.year * 10_000 + d.month * 100 + d.day


def _month_windows(start: date, end: date) -> list[MonthWindow]:
    """Return inclusive month windows from `start` to `end`."""
    out: list[MonthWindow] = []
    cur = date(start.year, start.month, 1)
    while cur <= end:
        if cur.month == 12:
            nxt = date(cur.year + 1, 1, 1)
        else:
            nxt = date(cur.year, cur.month + 1, 1)
        m_start = max(start, cur)
        m_end = min(end, nxt - timedelta(days=1))
        out.append(MonthWindow(m_start, m_end))
        cur = nxt
    return out


def _month_file_path(m: MonthWindow, root: str) -> Path:
    r = str(root).strip().upper()
    return OUT_DIR / f"{r.lower()}_1545_{m.start.year:04d}_{m.start.month:02d}.parquet"


def _trading_days_in_month(m: MonthWindow) -> pd.DatetimeIndex:
    """
    Approximate market sessions via business days.
    If you want exact NYSE holidays, replace with pandas-market-calendars.
    """
    return pd.bdate_range(m.start, m.end, freq="B")


def theta_get_json(
    session: requests.Session,
    path: str,
    params: dict[str, Any],
    *,
    log_http_errors: bool = True,
    allow_empty_472: bool = False,
    timeout_sec: float | None = None,
) -> dict[str, Any]:
    """
    GET JSON from the local Terminal with retries and timeouts.

    If ``allow_empty_472`` is True, HTTP 472 (no data) returns ``{}`` instead of
    raising. Raises ``requests.HTTPError`` on other HTTP errors after retries exhausted.
    """
    url = f"{THETA_BASE_URL.rstrip('/')}{path}"
    to = float(REQUEST_TIMEOUT_SEC if timeout_sec is None else timeout_sec)
    last_err: Exception | None = None
    for attempt in range(1, MAX_RETRIES + 1):
        try:
            r = session.get(
                url,
                params=params,
                headers=JSON_HEADERS,
                timeout=to,
            )
            if r.status_code == 472 and allow_empty_472:
                return {}
            r.raise_for_status()
            return r.json()
        except (requests.Timeout, requests.ConnectionError) as e:
            last_err = e
            if attempt == MAX_RETRIES:
                raise
            time.sleep(RETRY_SLEEP_SEC * attempt)
        except requests.HTTPError as e:
            if log_http_errors:
                _log(
                    f"HTTPError path={path} status={getattr(e.response, 'status_code', 'n/a')} "
                    f"attempt={attempt}/{MAX_RETRIES} params={params} "
                    f"body={_err_body(e.response)}"
                )
            # Do not retry most HTTP errors (4xx/5xx semantics vary).
            raise
    assert last_err is not None
    raise last_err


def _parse_tick_rows(
    payload: dict[str, Any],
    *,
    kind: str,
) -> pd.DataFrame:
    """
    Flatten Theta `response` list of {contract, ticks} into a DataFrame.

    `kind` is ``"quote"`` or ``"greeks"`` — uses documented tick column order
    when `header.format` is absent.
    """
    rows: list[dict[str, Any]] = []
    resp = payload.get("response")
    if not isinstance(resp, list):
        return pd.DataFrame()

    fmt = None
    hdr = payload.get("header")
    if isinstance(hdr, dict):
        fmt = hdr.get("format")

    for block in resp:
        if not isinstance(block, dict):
            continue
        c = block.get("contract") or {}
        ticks = block.get("ticks") or []
        if not isinstance(ticks, list):
            continue

        root = c.get("root")
        exp = c.get("expiration")
        strike = c.get("strike")
        right = c.get("right")

        for t in ticks:
            if not isinstance(t, (list, tuple)) or len(t) < 2:
                continue
            row: dict[str, Any] = {
                "root": root,
                "expiration": exp,
                "strike": strike,
                "right": right,
            }
            if fmt and isinstance(fmt, list):
                for i, name in enumerate(fmt):
                    if i < len(t):
                        row[str(name)] = t[i]
            elif kind == "quote":
                # ["ms_of_day","bid_size","bid_exchange","bid","bid_condition",
                #  "ask_size","ask_exchange","ask","ask_condition","date"]
                row["ms_of_day"] = t[0]
                row["bid"] = t[3] if len(t) > 3 else math.nan
                row["ask"] = t[7] if len(t) > 7 else math.nan
                row["date"] = t[-1]
            elif kind == "greeks":
                # ["ms_of_day","bid","ask","delta","theta","vega","rho","epsilon",
                #  "lambda","implied_vol","iv_error","ms_of_day2","underlying_price","date"]
                row["ms_of_day"] = t[0]
                row["bid"] = t[1] if len(t) > 1 else math.nan
                row["ask"] = t[2] if len(t) > 2 else math.nan
                row["delta"] = t[3] if len(t) > 3 else math.nan
                row["implied_vol"] = t[9] if len(t) > 9 else math.nan
                row["date"] = t[-1]
            rows.append(row)

    return pd.DataFrame(rows)


def _err_body(response: requests.Response | None) -> str:
    if response is None:
        return ""
    txt = (response.text or "").strip().replace("\n", " ")
    return txt[:400]


def _log(msg: str) -> None:
    ts = pd.Timestamp.now(tz=None).strftime("%Y-%m-%d %H:%M:%S")
    with LOG_FILE.open("a", encoding="utf-8") as f:
        f.write(f"[{ts}] {msg}\n")


def _fetch_greeks_by_expirations(
    session: requests.Session,
    day: date,
    q_df: pd.DataFrame,
    *,
    root: str,
    show_exp_progress: bool = True,
) -> pd.DataFrame:
    """
    Bulk greeks with ``exp=0`` does not return IV for the full chain; request
    greeks once per unique expiration present in the quote snapshot.
    """
    if q_df.empty or "expiration" not in q_df.columns:
        return pd.DataFrame()
    exps = (
        pd.to_numeric(q_df["expiration"], errors="coerce")
        .dropna()
        .astype(int)
        .unique()
        .tolist()
    )
    exps = sorted({int(x) for x in exps})
    ds = _yyyymmdd(day)
    chunks: list[pd.DataFrame] = []
    exp_iter: Iterable[int] = exps
    if show_exp_progress and len(exps) > 1:
        exp_iter = tqdm(
            exps,
            desc="greeks",
            leave=False,
            unit="exp",
            mininterval=0.15,
        )
    for exp in exp_iter:
        greek_params: dict[str, Any] = {
            "root": str(root).strip().upper(),
            "exp": exp,
            "start_date": ds,
            "end_date": ds,
            "ivl": GREEKS_IVL_MS,
            "start_time": str(GREEKS_WINDOW_START_MS),
            "end_time": str(GREEKS_WINDOW_END_MS),
            "rth": True,
            "use_csv": False,
            "pretty_time": False,
        }
        g_json = theta_get_json(
            session,
            "/v2/bulk_hist/option/greeks",
            greek_params,
            log_http_errors=False,
            allow_empty_472=True,
        )
        part = _parse_tick_rows(g_json, kind="greeks")
        if not part.empty:
            chunks.append(part)
        time.sleep(GREEKS_INTER_EXP_SLEEP_SEC)
    if not chunks:
        return pd.DataFrame()
    full = pd.concat(chunks, ignore_index=True)
    return _nearest_time_tick(full, AT_TIME_MS)


def _nearest_time_tick(df: pd.DataFrame, target_ms: int) -> pd.DataFrame:
    """Keep one row per (expiration, strike, right) with ms_of_day closest to target_ms."""
    if df.empty or "ms_of_day" not in df.columns:
        return df
    work = df.copy()
    work["ms_of_day"] = pd.to_numeric(work["ms_of_day"], errors="coerce")
    work["__dt__"] = (work["ms_of_day"] - float(target_ms)).abs()
    return (
        work.sort_values("__dt__")
        .groupby(["expiration", "strike", "right"], as_index=False)
        .first()
        .drop(columns=["__dt__"], errors="ignore")
    )


def _merge_quote_greeks(quote_df: pd.DataFrame, greeks_df: pd.DataFrame) -> pd.DataFrame:
    """Merge NBBO (quote) with IV/delta (greeks) on contract + session date."""
    if quote_df.empty:
        return pd.DataFrame()
    if greeks_df.empty:
        return quote_df

    keys = ["expiration", "strike", "right", "date"]
    g_cols = [k for k in keys if k in greeks_df.columns]
    extra = [c for c in ("implied_vol", "delta") if c in greeks_df.columns]
    g = greeks_df[g_cols + extra].drop_duplicates()
    return quote_df.merge(g, on=keys, how="left")


def _normalize_to_schema(
    df: pd.DataFrame,
    *,
    allow_quote_only: bool = False,
) -> pd.DataFrame:
    """
    Map merged DataFrame to:
    [quote_datetime, expiration, strike, right, bid, ask, implied_vol, delta]

    If ``allow_quote_only`` is True, rows without vendor ``implied_vol`` are kept
    (NaN greeks) so you can backfill IV/delta later (e.g. BS inversion).
    """
    if df.empty:
        return pd.DataFrame(
            columns=[
                "quote_datetime",
                "expiration",
                "strike",
                "right",
                "bid",
                "ask",
                "implied_vol",
                "delta",
            ]
        )

    exp_raw = pd.to_numeric(df["expiration"], errors="coerce").astype("Int64")
    strike_raw = pd.to_numeric(df["strike"], errors="coerce")
    out = pd.DataFrame(
        {
            "expiration": pd.to_datetime(exp_raw.astype(str), format="%Y%m%d", errors="coerce").dt.normalize(),
            "strike": (strike_raw / 10_000.0).astype(float),
            "right": df["right"].astype(str).str.upper().str[0],
            "bid": pd.to_numeric(df.get("bid"), errors="coerce"),
            "ask": pd.to_numeric(df.get("ask"), errors="coerce"),
            "implied_vol": pd.to_numeric(df.get("implied_vol"), errors="coerce"),
            "delta": pd.to_numeric(df.get("delta"), errors="coerce"),
        }
    )

    d_int = pd.to_numeric(df["date"], errors="coerce").astype("Int64")
    ms = pd.to_numeric(df.get("ms_of_day"), errors="coerce")
    # Build ET timestamp: date + ms since midnight
    base = pd.to_datetime(d_int.astype(str), format="%Y%m%d", errors="coerce")
    qdt = base + pd.to_timedelta(ms, unit="ms")
    out["quote_datetime"] = qdt.dt.tz_localize(
        "America/New_York", nonexistent="shift_forward", ambiguous="NaT"
    )

    keep = (
        out["quote_datetime"].notna()
        & out["expiration"].notna()
        & (out["strike"] > 0)
        & out["right"].isin(["C", "P"])
        & (out["bid"] > 0)
        & (out["ask"] > 0)
    )
    if not allow_quote_only:
        keep = keep & out["implied_vol"].notna()
    out = out.loc[keep].copy()

    out["strike"] = out["strike"].astype("float32")
    out["bid"] = out["bid"].astype("float32")
    out["ask"] = out["ask"].astype("float32")
    out["implied_vol"] = out["implied_vol"].astype("float32")
    out["delta"] = out["delta"].astype("float32")
    out["right"] = out["right"].astype("category")
    return out.reset_index(drop=True)


def fetch_one_day(
    session: requests.Session,
    day: date,
    *,
    root: str = DEFAULT_ROOT,
    show_greeks_exp_progress: bool = True,
    allow_quote_only: bool = False,
    skip_greeks: bool = False,
) -> pd.DataFrame:
    """
    Pull full option chain at ~15:45 ET for `day` via REST, return normalized schema rows.
    """
    sym = str(root).strip().upper()
    ds = _yyyymmdd(day)

    # 1) NBBO at exact clock time (bulk_at_time / quote).
    quote_params: dict[str, Any] = {
        "root": sym,
        "exp": 0,
        "start_date": ds,
        "end_date": ds,
        "ivl": AT_TIME_MS,
        "rth": True,
        "use_csv": False,
        "pretty_time": False,
    }
    # 472 = Theta "no data" (e.g. exchange holiday, early close, gap in history).
    q_json = theta_get_json(
        session,
        "/v2/bulk_at_time/option/quote",
        quote_params,
        allow_empty_472=True,
    )
    q_df = _parse_tick_rows(q_json, kind="quote")

    # 2) Greeks / IV: one request per expiration (exp=0 is invalid for this endpoint).
    if skip_greeks:
        g_df = pd.DataFrame()
    else:
        g_df = _fetch_greeks_by_expirations(
            session,
            day,
            q_df,
            root=sym,
            show_exp_progress=show_greeks_exp_progress,
        )

    merged = _merge_quote_greeks(q_df, g_df)
    quote_only = allow_quote_only or skip_greeks
    out = _normalize_to_schema(merged, allow_quote_only=quote_only)
    _log(
        f"day={day} quote_rows={len(q_df)} greeks_rows={len(g_df)} "
        f"skip_greeks={skip_greeks} merged_rows={len(merged)} final_rows={len(out)}"
    )
    return out


def _fetch_month_data(
    session: requests.Session,
    m: MonthWindow,
    *,
    root: str = DEFAULT_ROOT,
    show_progress: bool = True,
    allow_quote_only: bool = False,
    skip_greeks: bool = False,
) -> pd.DataFrame:
    """
    Fetch one month of 15:45 ET snapshots (daily REST calls), with retries per day.
    """
    month_rows: list[pd.DataFrame] = []
    trading_days = _trading_days_in_month(m)
    ym = f"{m.start.year}-{m.start.month:02d}"

    day_iter: Iterable[pd.Timestamp] = trading_days
    if show_progress:
        day_iter = tqdm(
            trading_days,
            desc=f"{ym} days",
            leave=False,
            unit="day",
        )

    for ts in day_iter:
        day = ts.date()
        for attempt in range(1, MAX_RETRIES + 1):
            try:
                ddf = fetch_one_day(
                    session,
                    day,
                    root=root,
                    show_greeks_exp_progress=show_progress,
                    allow_quote_only=allow_quote_only,
                    skip_greeks=skip_greeks,
                )
                if not ddf.empty:
                    month_rows.append(ddf)
                break
            except (requests.Timeout, requests.ConnectionError, requests.HTTPError) as e:
                if attempt == MAX_RETRIES:
                    _log(f"day_failed day={day} err={type(e).__name__}: {e}")
                    if show_progress:
                        print(f"[WARN] Day {day} failed after retries: {e}")
                else:
                    time.sleep(RETRY_SLEEP_SEC * attempt)

    if not month_rows:
        return pd.DataFrame(
            columns=[
                "quote_datetime",
                "expiration",
                "strike",
                "right",
                "bid",
                "ask",
                "implied_vol",
                "delta",
            ]
        )
    return pd.concat(month_rows, ignore_index=True)


def _month_worker_job(
    payload: tuple[MonthWindow, bool, float, bool, str],
) -> tuple[str, int, str | None]:
    """
    Download one calendar month and write its parquet file.

    Intended for ``ProcessPoolExecutor`` (separate process, own HTTP session).
    ``payload`` is ``(month_window, allow_quote_only, request_timeout_sec, skip_greeks, root)``.
    Returns ``(parquet_filename, row_count, error_message)``.
    """
    global REQUEST_TIMEOUT_SEC
    m, allow_quote_only, request_timeout_sec, skip_greeks, root = payload
    REQUEST_TIMEOUT_SEC = float(request_timeout_sec)
    global _WORKER_GREEKS_WINDOW_LOGGED
    if not skip_greeks and not _WORKER_GREEKS_WINDOW_LOGGED:
        _WORKER_GREEKS_WINDOW_LOGGED = True
        _log(
            f"worker_process greeks_ivl_ms={GREEKS_IVL_MS} "
            f"start_time={GREEKS_WINDOW_START_MS} end_time={GREEKS_WINDOW_END_MS}"
        )
    out_file = _month_file_path(m, root)
    session = requests.Session()
    try:
        with session:
            month_df = _fetch_month_data(
                session,
                m,
                root=root,
                show_progress=False,
                allow_quote_only=allow_quote_only,
                skip_greeks=skip_greeks,
            )
            month_df.to_parquet(out_file, index=False, compression="zstd")
        _log(f"month_done file={out_file.name} rows={len(month_df)}")
        return (out_file.name, len(month_df), None)
    except Exception as e:  # noqa: BLE001
        _log(f"month_fail ym={m.start.year:04d}-{m.start.month:02d} err={e}")
        return (out_file.name, 0, str(e))
    finally:
        gc.collect()


def _months_to_process(months: list[MonthWindow], root: str) -> list[MonthWindow]:
    """Months that still need a pull (missing file, or existing file with 0 rows)."""
    todo: list[MonthWindow] = []
    for m in months:
        out_file = _month_file_path(m, root)
        if out_file.exists():
            existing_rows = _parquet_num_rows(out_file)
            if existing_rows > 0:
                continue
            _log(
                f"Re-fetching empty month file file={out_file.name} rows={existing_rows}"
            )
        todo.append(m)
    return todo


def build_dataset(
    *,
    root: str = DEFAULT_ROOT,
    start_date: date | None = None,
    end_date: date | None = None,
    workers: int = 1,
    allow_quote_only: bool = False,
    skip_greeks: bool = False,
    request_timeout_sec: float = REQUEST_TIMEOUT_SEC,
) -> None:
    """Main orchestrator: month-by-month fetch, parquet write, resume skipping."""
    global REQUEST_TIMEOUT_SEC
    request_timeout_sec = float(request_timeout_sec)
    if request_timeout_sec < 10:
        raise ValueError("request_timeout_sec must be >= 10")
    REQUEST_TIMEOUT_SEC = request_timeout_sec

    sym = str(root).strip().upper()
    sd = start_date if start_date is not None else START_DATE
    ed = end_date if end_date is not None else END_DATE
    if ed < sd:
        raise ValueError("end_date must be on or after start_date")

    months = _month_windows(sd, ed)
    print(
        f"Starting ThetaData {sym} 15:45 pull (REST): {sd} -> {ed} "
        f"({len(months)} months)"
    )
    print(f"Terminal base: {THETA_BASE_URL}")
    print(f"Output dir: {OUT_DIR}")
    print(f"Debug log: {LOG_FILE}")
    print(f"Workers: {workers}", flush=True)
    print(
        f"Allow quote-only rows (missing IV/delta): {allow_quote_only}",
        flush=True,
    )
    print(f"Skip greeks endpoint (quote-only pull): {skip_greeks}", flush=True)
    print(f"HTTP read timeout (seconds): {REQUEST_TIMEOUT_SEC}", flush=True)
    if skip_greeks:
        print(
            "Progress: per-day quote pulls only (no per-expiration greeks). "
            f"With --workers > 1, child processes log to {LOG_FILE.name}.",
            flush=True,
        )
    else:
        print(
            "Progress: the top bar advances once per month. While a month runs you will see "
            "a per-day bar and a per-expiration greeks bar (dozens of HTTP calls per day). "
            "With --workers > 1, only the month-level bar prints here; child processes log to "
            f"{LOG_FILE.name}.",
            flush=True,
        )

    months_to_do = _months_to_process(months, sym)
    if not months_to_do:
        print("All month files already exist with data; nothing to do.", flush=True)
        return

    _log(
        f"START run root={sym} start={sd} end={ed} "
        f"calendar_months={len(months)} months_to_do={len(months_to_do)} "
        f"workers={workers} base={THETA_BASE_URL} "
        f"at_time_ms={AT_TIME_MS} greeks_ivl_ms={GREEKS_IVL_MS} "
        f"greeks_start_time={GREEKS_WINDOW_START_MS} greeks_end_time={GREEKS_WINDOW_END_MS} "
        f"allow_quote_only={allow_quote_only} skip_greeks={skip_greeks} "
        f"request_timeout_sec={REQUEST_TIMEOUT_SEC}"
    )
    if not skip_greeks:
        print(
            f"Greeks window (Option A): ivl_ms={GREEKS_IVL_MS}, "
            f"start_time={GREEKS_WINDOW_START_MS}, end_time={GREEKS_WINDOW_END_MS} "
            f"(span must exceed ivl; 15:44–~15:47 ET around 15:45 snapshot).",
            flush=True,
        )

    if workers < 1:
        raise ValueError("workers must be >= 1")

    if workers == 1:
        session = requests.Session()
        with session:
            for m in tqdm(months_to_do, desc="Processing months", unit="month"):
                out_file = _month_file_path(m, sym)
                try:
                    month_df = _fetch_month_data(
                        session,
                        m,
                        root=sym,
                        show_progress=True,
                        allow_quote_only=allow_quote_only,
                        skip_greeks=skip_greeks,
                    )
                    month_df.to_parquet(out_file, index=False, compression="zstd")
                    _log(f"month_done file={out_file.name} rows={len(month_df)}")
                except Exception as e:  # noqa: BLE001
                    _log(f"month_fail ym={m.start.year:04d}-{m.start.month:02d} err={e}")
                    raise RuntimeError(
                        f"Failed on month {m.start.year:04d}-{m.start.month:02d}: {e}"
                    ) from e
                finally:
                    if "month_df" in locals():
                        del month_df
                    gc.collect()
    else:
        with ProcessPoolExecutor(max_workers=workers) as pool:
            future_map = {
                pool.submit(
                    _month_worker_job,
                    (m, allow_quote_only, REQUEST_TIMEOUT_SEC, skip_greeks, sym),
                ): m
                for m in months_to_do
            }
            for fut in tqdm(
                as_completed(future_map),
                total=len(future_map),
                desc="Processing months",
                unit="month",
            ):
                name, rows, err = fut.result()
                if err is not None:
                    raise RuntimeError(f"Failed on month file {name}: {err}") from None

    print("Done.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Download 15:45 ET option snapshots from local Theta Terminal (default: SPY)."
    )
    parser.add_argument(
        "--root",
        type=str,
        default=DEFAULT_ROOT,
        help="Option root symbol (default SPY). Examples: VXX, VIX.",
    )
    parser.add_argument(
        "--start-date",
        type=str,
        default=None,
        help="First calendar date YYYY-MM-DD (default: built-in START_DATE).",
    )
    parser.add_argument(
        "--end-date",
        type=str,
        default=None,
        help="Last calendar date YYYY-MM-DD (default: today).",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=1,
        help="Parallel month downloads (default 1). Try 2–4; higher may overload the Terminal.",
    )
    parser.add_argument(
        "--allow-quote-only",
        action="store_true",
        help=(
            "Keep rows with valid bid/ask even when vendor implied_vol/delta are missing "
            "(NaN greeks for later BS / enrichment)."
        ),
    )
    parser.add_argument(
        "--skip-greeks",
        action="store_true",
        help=(
            "Do not call /v2/bulk_hist/option/greeks; output quote-only rows with NaN "
            "implied_vol/delta (same filter as --allow-quote-only for IV)."
        ),
    )
    parser.add_argument(
        "--timeout-sec",
        type=float,
        default=300.0,
        help=(
            "HTTP read timeout per request to the local Terminal (default 300). "
            "Increase if you see ReadTimeout on large greeks responses."
        ),
    )
    args = parser.parse_args()
    sd = date.fromisoformat(args.start_date) if args.start_date else None
    ed = date.fromisoformat(args.end_date) if args.end_date else None
    build_dataset(
        root=str(args.root),
        start_date=sd,
        end_date=ed,
        workers=max(1, args.workers),
        allow_quote_only=bool(args.allow_quote_only),
        skip_greeks=bool(args.skip_greeks),
        request_timeout_sec=float(args.timeout_sec),
    )