"""
Intraday bars from Alpaca's market-data API (free Basic plan: full-market SIP
bars from 2016, all but the latest 15 minutes; delisted symbols included).

    PYTHONPATH=. uv run python -m research.lab.fetch_alpaca etfs              # 1-min SPY QQQ IWM DIA TLT GLD
    PYTHONPATH=. uv run python -m research.lab.fetch_alpaca stocks [--universe us|midsmall|all]   # 5-min PIT ever-members
    PYTHONPATH=. uv run python -m research.lab.fetch_alpaca status

Keys come from ALPACA_API_KEY / ALPACA_API_SECRET in the environment or .env
and are never printed. Bars are split- and dividend-adjusted
(adjustment=all), regular session only (09:30-16:00 New York; a bar is
labeled by its start time), stored per symbol as
data/research/intraday/alpaca/<timeframe>/<SYMBOL>.pkl, float32 columns
o h l c v vw n, index tz-aware New York. A symbol's file is written once
complete, so a rerun resumes. An empty result is recorded in _empty.json and
not retried.
"""

from __future__ import annotations

import argparse
import json
import os
import threading
import time
from concurrent.futures import ThreadPoolExecutor

import numpy as np
import pandas as pd
import requests

ROOT = os.path.join("data", "research", "intraday", "alpaca")
URL = "https://data.alpaca.markets/v2/stocks/{sym}/bars"
START, END = "2016-01-01", None
ETFS = ["SPY", "QQQ", "IWM", "DIA", "TLT", "GLD"]
MIN_INTERVAL = 60 / 180                     # stay under the 200 requests/minute limit
WORKERS = 6                                 # requests take ~1-2 s; six threads reach the limit
NY = "America/New_York"


def _keys() -> dict[str, str]:
    env = dict(os.environ)
    if os.path.exists(".env"):
        for line in open(".env"):
            if "=" in line and not line.lstrip().startswith("#"):
                k, v = line.split("=", 1)
                env.setdefault(k.strip(), v.strip().strip("'\""))
    return {"APCA-API-KEY-ID": env["ALPACA_API_KEY"], "APCA-API-SECRET-KEY": env["ALPACA_API_SECRET"]}


class Client:
    """Thread-safe: one session per thread, one shared request clock (<= 180/min)."""

    def __init__(self) -> None:
        self.headers = _keys()
        self.local = threading.local()
        self.lock = threading.Lock()
        self.next_slot = 0.0

    def _session(self) -> requests.Session:
        if not hasattr(self.local, "s"):
            self.local.s = requests.Session()
            self.local.s.headers.update(self.headers)
        return self.local.s

    def _wait_turn(self) -> None:
        with self.lock:
            now = time.time()
            slot = max(now, self.next_slot)
            self.next_slot = slot + MIN_INTERVAL
        if slot > now:
            time.sleep(slot - now)

    def get(self, sym: str, params: dict) -> dict:
        for attempt in range(8):
            self._wait_turn()
            try:
                r = self._session().get(URL.format(sym=sym), params=params, timeout=60)
            except requests.RequestException:
                time.sleep(2 ** attempt)
                continue
            if r.status_code == 429 or r.status_code >= 500:
                time.sleep(2 ** attempt)
                continue
            if r.status_code != 200:
                raise RuntimeError(f"{sym}: HTTP {r.status_code} {r.text[:200]}")
            return r.json()
        raise RuntimeError(f"{sym}: retries exhausted")


def fetch_symbol(cl: Client, sym: str, timeframe: str, start: str, end: str) -> pd.DataFrame:
    params = {"timeframe": timeframe, "start": start, "end": end, "feed": "sip", "limit": 10000,
              "adjustment": "all"}
    parts = []
    while True:
        j = cl.get(sym.replace("-", "."), params)
        if j.get("bars"):
            df = pd.DataFrame(j["bars"])
            df.index = pd.DatetimeIndex(pd.to_datetime(df.pop("t"), utc=True)).tz_convert(NY)
            hm = df.index.hour * 60 + df.index.minute
            df = df[(hm >= 9 * 60 + 30) & (hm < 16 * 60)]
            parts.append(df.reindex(columns=["o", "h", "l", "c", "v", "vw", "n"]).astype(np.float32))
        tok = j.get("next_page_token")
        if not tok:
            break
        params["page_token"] = tok
    return pd.concat(parts) if parts else pd.DataFrame()


def run(symbols: list[str], timeframe: str) -> None:
    out = os.path.join(ROOT, timeframe)
    os.makedirs(out, exist_ok=True)
    empty_path = os.path.join(out, "_empty.json")
    empty = set(json.load(open(empty_path))) if os.path.exists(empty_path) else set()
    end = (pd.Timestamp.now(tz="UTC") - pd.Timedelta(minutes=20)).strftime("%Y-%m-%dT%H:%M:%SZ")
    cl = Client()
    todo = [s for s in symbols if s not in empty and not os.path.exists(os.path.join(out, f"{s}.pkl"))]
    print(f"{timeframe}: {len(symbols)} symbols, {len(todo)} to fetch", flush=True)
    t0 = time.time()
    lock = threading.Lock()
    done = [0]

    def one(s: str) -> None:
        try:
            df = fetch_symbol(cl, s, timeframe, START, end)
        except RuntimeError as e:
            print(f"  {s}: {e}", flush=True)
            return
        if not df.empty:
            tmp = os.path.join(out, f".{s}.pkl.tmp")
            df.to_pickle(tmp)
            os.replace(tmp, os.path.join(out, f"{s}.pkl"))
        with lock:
            if df.empty:
                empty.add(s)
                json.dump(sorted(empty), open(empty_path, "w"))
            done[0] += 1
            i = done[0]
        if i % 10 == 0 or i == len(todo):
            print(f"  {i}/{len(todo)} {s} rows {len(df)} elapsed {(time.time()-t0)/60:.1f} min", flush=True)

    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        list(ex.map(one, todo))


def _liquid(p) -> pd.DataFrame:
    """Member days that pass the day-trading filters of research/lab/orb_stocks_in_play.py
    (14-day mean volume >= 1M shares, close > $5, 14-day mean range > $0.50) on daily bars."""
    vol = p.volume.rolling(14, min_periods=10).mean()
    rng = (p.high - p.low).rolling(14, min_periods=10).mean()
    return p.member & (vol >= 1e6) & (p.close > 5) & (rng > 0.5)


def universe(name: str, liquid: bool = True) -> list[str]:
    """PIT members at any time since 2016 (Yahoo spelling); with ``liquid``, only names
    that passed the day-trading filters on at least one member day."""
    syms: set[str] = set()
    panels = []
    if name in ("us", "all"):
        from research.lab.factor_scan import load
        panels.append(load("us")[0])
    if name in ("midsmall", "all"):
        from research.lab.panel_mid import load_us_mid
        panels.append(load_us_mid("midsmall"))
    for p in panels:
        m = (_liquid(p) if liquid else p.member).loc[START:]
        syms |= set(m.columns[m.any()])
    return sorted(s for s in syms if s.replace("-", "").replace(".", "").isalnum())


def status() -> None:
    for tf in sorted(os.listdir(ROOT)) if os.path.exists(ROOT) else []:
        d = os.path.join(ROOT, tf)
        files = [f for f in os.listdir(d) if f.endswith(".pkl")]
        empty = json.load(open(os.path.join(d, "_empty.json"))) if os.path.exists(os.path.join(d, "_empty.json")) else []
        size = sum(os.path.getsize(os.path.join(d, f)) for f in files) / 1e9
        print(f"{tf}: {len(files)} symbols, {len(empty)} empty, {size:.2f} GB")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("what", choices=["etfs", "stocks", "status"])
    ap.add_argument("--universe", default="all", choices=["us", "midsmall", "all"])
    ap.add_argument("--all-members", action="store_true", help="skip the liquidity pre-filter")
    a = ap.parse_args()
    if a.what == "status":
        status()
    elif a.what == "etfs":
        run(ETFS, "1Min")
    else:
        run(universe(a.universe, liquid=not a.all_members), "5Min")


if __name__ == "__main__":
    main()