"""
Fetch a long OHLCV history for every ticker that was ever an S&P 500 member.

The live panels in data/us.csv hold 10 years of close/open for tickers with
>50% coverage. Factor research needs more: volume and high/low, 20+ years of
history for an out-of-sample design split, and recently listed members (CEG,
GEV, ...) that the coverage rule drops. Delisted names that Yahoo no longer
serves stay missing; ``coverage_report`` in panel.py quantifies that gap.

Usage:
    uv run python -m research.lab.fetch_us [--start 2004-01-01]

Writes data/research/us_ohlcv.pkl: {"open","high","low","close","volume"}
DataFrames (dates x tickers), split/dividend adjusted.
"""

import argparse
import os
import pickle
import time

import pandas as pd
import yfinance as yf

import universe_history as uh

OUT = os.path.join("data", "research", "us_ohlcv.pkl")
FIELDS = ["Open", "High", "Low", "Close", "Volume"]
# Benchmarks, cash proxies and sector ETFs used by overlays and attribution.
EXTRA = ["SPY", "RSP", "QQQ", "SHY", "BIL", "IEF", "TLT", "GLD",
         "XLB", "XLE", "XLF", "XLI", "XLK", "XLP", "XLU", "XLV", "XLY"]


def _proxy_env() -> None:
    proxy = os.environ.get("QUANT_MARKET_PROXY")
    if not proxy and os.path.exists(".env"):
        for line in open(".env"):
            if line.startswith("QUANT_MARKET_PROXY="):
                proxy = line.split("=", 1)[1].strip().strip("'\"")
    if proxy:
        for k in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
            os.environ[k] = proxy


def fetch(tickers: list[str], start: str, chunk: int = 60) -> dict[str, pd.DataFrame]:
    parts: dict[str, list[pd.DataFrame]] = {f: [] for f in FIELDS}
    for i in range(0, len(tickers), chunk):
        batch = tickers[i:i + chunk]
        for attempt in range(4):
            try:
                raw = yf.download(batch, start=start, auto_adjust=True, progress=False,
                                  group_by="column", threads=True)
                break
            except Exception as exc:  # noqa: BLE001 — retry transient Yahoo errors
                print(f"  batch {i}: {exc}; retry {attempt + 1}")
                time.sleep(10 * (attempt + 1))
        else:
            continue
        for f in FIELDS:
            if f in raw.columns.get_level_values(0):
                parts[f].append(raw[f])
        print(f"  fetched {min(i + chunk, len(tickers))}/{len(tickers)}", flush=True)
        time.sleep(2)
    out = {}
    for f in FIELDS:
        if not parts[f]:
            raise RuntimeError(f"no {f} data downloaded")
        df = pd.concat(parts[f], axis=1)
        df = df.loc[:, ~df.columns.duplicated()].sort_index()
        out[f.lower()] = df.dropna(axis=1, how="all")
    return out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--start", default="2004-01-01")
    args = ap.parse_args()
    _proxy_env()
    intervals = uh.load_sp500_history()
    tickers = sorted(set(uh.all_tickers_ever(intervals)) | set(EXTRA))
    print(f"Fetching {len(tickers)} tickers from {args.start}")
    panel = fetch(tickers, args.start)
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "wb") as fh:
        pickle.dump(panel, fh)
    c = panel["close"]
    print(f"Saved {OUT}: {c.shape[1]} tickers, {c.index[0].date()} → {c.index[-1].date()}")


if __name__ == "__main__":
    main()