#!/usr/bin/env python3
"""
Find NYSE (XNYS) session dates missing from ``spy_1545_YYYY_MM.parquet`` chunks and
refetch them via the local Theta Terminal (same HTTP flow as ``build_theta_dataset``).

Scope: only calendar months that already have a chunk file under ``theta_chunks/``.
Install: ``pip install exchange_calendars`` (optional fallback uses weekdays minus a
static US holiday list — less accurate).

Example::

    python RenTech/data_pipeline/backfill_theta_missing_days.py --dry-run
    python RenTech/data_pipeline/backfill_theta_missing_days.py --skip-greeks
"""

from __future__ import annotations

import argparse
import sys
from datetime import date, datetime, timedelta
from pathlib import Path

import pandas as pd
import requests
from tqdm import tqdm

_DATA_PIPE = Path(__file__).resolve().parent
if str(_DATA_PIPE) not in sys.path:
    sys.path.insert(0, str(_DATA_PIPE))

import build_theta_dataset as btd  # noqa: E402

# Static holidays (approximation if exchange_calendars missing). NYSE-only subset.
_US_HOLIDAYS_FALLBACK: set[tuple[int, int]] = {
    (1, 1),
    (7, 4),
    (12, 25),
}


def _xnys_dates_in_range(lo: date, hi: date) -> set[date]:
    try:
        import exchange_calendars as xcals

        cal = xcals.get_calendar("XNYS")
        # Use naive calendar dates; avoid sessions_in_range + tz=UTC (pandas 2.x can use
        # datetime.timezone.utc and break exchange_calendars' parse_date).
        out: set[date] = set()
        d = lo
        while d <= hi:
            if cal.is_session(pd.Timestamp(d)):
                out.add(d)
            d += timedelta(days=1)
        return out
    except ImportError:
        # Weekdays only; misses Good Friday, Juneteenth, etc.
        d = lo
        out: set[date] = set()
        while d <= hi:
            if d.weekday() < 5 and (d.month, d.day) not in _US_HOLIDAYS_FALLBACK:
                out.add(d)
            d += timedelta(days=1)
        return out


def _month_bounds(y: int, m: int) -> tuple[date, date]:
    first = date(y, m, 1)
    if m == 12:
        last = date(y + 1, 1, 1) - timedelta(days=1)
    else:
        last = date(y, m + 1, 1) - timedelta(days=1)
    return first, last


def _parse_chunk_name(path: Path) -> tuple[str, int, int] | None:
    """
    Parse ``{ROOT}_1545_YYYY_MM`` (e.g. ``spy_1545_2020_01``, ``vxx_1545_2020_01``).
    """
    stem = path.stem
    if "_1545_" not in stem:
        return None
    root_part, rest = stem.split("_1545_", 1)
    parts = rest.split("_")
    if len(parts) != 2:
        return None
    try:
        return root_part.upper(), int(parts[0]), int(parts[1])
    except ValueError:
        return None


def _quote_session_dates_in_parquet(path: Path) -> set[date]:
    df = pd.read_parquet(path, columns=["quote_datetime"])
    if df.empty:
        return set()
    ts = pd.to_datetime(df["quote_datetime"])
    if ts.dt.tz is None:
        ts = ts.dt.tz_localize(
            "America/New_York",
            ambiguous="NaT",
            nonexistent="shift_forward",
        )
    else:
        ts = ts.dt.tz_convert("America/New_York")
    ts = ts[ts.notna()]
    return set(ts.dt.date.unique())


def _merge_day_into_month_file(month_path: Path, day: date, new_rows: pd.DataFrame) -> None:
    if new_rows.empty:
        return
    if not month_path.exists():
        new_rows.sort_values("quote_datetime").reset_index(drop=True).to_parquet(
            month_path, index=False, compression="zstd"
        )
        return

    existing = pd.read_parquet(month_path)
    ts = pd.to_datetime(existing["quote_datetime"])
    if ts.dt.tz is None:
        ts = ts.dt.tz_localize(
            "America/New_York",
            ambiguous="NaT",
            nonexistent="shift_forward",
        )
    else:
        ts = ts.dt.tz_convert("America/New_York")
    keep = ts.dt.date != day
    trimmed = existing.loc[keep].copy()
    out = pd.concat([trimmed, new_rows], ignore_index=True)
    out = out.sort_values("quote_datetime").reset_index(drop=True)
    out.to_parquet(month_path, index=False, compression="zstd")


def collect_missing_by_month(
    out_dir: Path,
    *,
    max_session_date: date | None = None,
) -> tuple[list[tuple[Path, date, set[date]]], bool]:
    """
    Returns list of (month_path, month_first, missing_dates) and whether XNYS calendar was used.
    """
    try:
        import exchange_calendars as xcals  # noqa: F401

        used_xnys = True
    except ImportError:
        used_xnys = False

    chunks = sorted(out_dir.glob("*_1545_*.parquet"))
    report: list[tuple[Path, date, set[date]]] = []

    for path in chunks:
        ym = _parse_chunk_name(path)
        if ym is None:
            continue
        _root, y, m = ym
        first, last = _month_bounds(y, m)
        hi = last
        if max_session_date is not None:
            hi = min(last, max_session_date)
        if hi < first:
            continue
        expected = _xnys_dates_in_range(first, hi)
        present = _quote_session_dates_in_parquet(path)
        missing = sorted(expected - present)
        if missing:
            report.append((path, first, set(missing)))

    return report, used_xnys


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill NYSE session dates missing from Theta SPY 15:45 parquet months.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Only list missing dates; do not call the Terminal or write files.",
    )
    parser.add_argument(
        "--skip-greeks",
        action="store_true",
        help="Same as build_theta_dataset: quotes only.",
    )
    parser.add_argument(
        "--allow-quote-only",
        action="store_true",
        help="Pass through when also fetching greeks (usually redundant with --skip-greeks).",
    )
    parser.add_argument(
        "--timeout-sec",
        type=float,
        default=btd.REQUEST_TIMEOUT_SEC,
        help=f"HTTP timeout (default {btd.REQUEST_TIMEOUT_SEC}).",
    )
    parser.add_argument(
        "--max-date",
        type=str,
        default="",
        help="YYYY-MM-DD: only sessions on or before this date (default: today).",
    )
    parser.add_argument(
        "--include-future-sessions",
        action="store_true",
        help="Do not cap at today (refetch missing XNYS dates even if after calendar today).",
    )
    args = parser.parse_args()

    max_session_date: date | None
    if args.include_future_sessions:
        max_session_date = None
    elif args.max_date.strip():
        max_session_date = datetime.strptime(args.max_date.strip(), "%Y-%m-%d").date()
    else:
        max_session_date = date.today()

    out_dir = btd.OUT_DIR
    report, used_xnys = collect_missing_by_month(out_dir, max_session_date=max_session_date)

    all_missing: list[tuple[Path, date]] = []
    for month_path, _first, dates in report:
        for d in sorted(dates):
            all_missing.append((month_path, d))

    if not used_xnys:
        print(
            "[WARN] exchange_calendars not installed; using weekday fallback (install for "
            "accurate NYSE holidays). pip install exchange_calendars",
            flush=True,
        )
    print(f"Chunk dir: {out_dir}", flush=True)
    print(
        f"Max session date (inclusive): {max_session_date if max_session_date else 'none (all future allowed)'}",
        flush=True,
    )
    print(f"Months with gaps: {len(report)}", flush=True)
    print(f"Missing session-days to fetch: {len(all_missing)}", flush=True)

    if not all_missing:
        print("Nothing missing.", flush=True)
        return

    for month_path, d in sorted(all_missing, key=lambda x: (x[1], x[0].name))[:50]:
        print(f"  {d}  ->  {month_path.name}", flush=True)
    if len(all_missing) > 50:
        print(f"  ... and {len(all_missing) - 50} more", flush=True)

    if args.dry_run:
        return

    btd.REQUEST_TIMEOUT_SEC = float(args.timeout_sec)
    session = requests.Session()
    ok = 0
    empty = 0
    err = 0

    for month_path, day in tqdm(all_missing, desc="Backfilling days", unit="day"):
        parsed = _parse_chunk_name(month_path)
        opt_root = parsed[0] if parsed else btd.DEFAULT_ROOT
        try:
            df = btd.fetch_one_day(
                session,
                day,
                root=opt_root,
                show_greeks_exp_progress=False,
                allow_quote_only=bool(args.allow_quote_only),
                skip_greeks=bool(args.skip_greeks),
            )
            if df.empty:
                btd._log(f"backfill_empty day={day} file={month_path.name}")
                empty += 1
                continue
            _merge_day_into_month_file(month_path, day, df)
            btd._log(f"backfill_ok day={day} rows={len(df)} file={month_path.name}")
            ok += 1
        except Exception as e:  # noqa: BLE001
            btd._log(f"backfill_fail day={day} file={month_path.name} err={e}")
            err += 1
            print(f"[ERROR] {day} {month_path.name}: {e}", flush=True)

    print(f"Done. wrote_rows: {ok}  still_empty_from_theta: {empty}  errors: {err}", flush=True)
    session.close()


if __name__ == "__main__":
    main()