#!/usr/bin/env python3
"""
Download Alpaca 1-minute RTH bars and **append** to ``SP500_1Min_Parquet_RTH_FULL``.

Reads API keys from ``trading_bot_live/.env`` (``APCA_API_KEY_ID``, ``APCA_API_SECRET_KEY``).
Timestamps stored as **naive UTC** (09:30 ET → 14:30 UTC in winter), matching existing parquets.

Example (smoke)::

    cd /Users/robzingale/trading_bot && PYTHONUNBUFFERED=1 .venv/bin/python \\
        RenTech/data_pipeline/download_alpaca_rth_minutes.py \\
        --start 2025-01-02 --symbols SPY,AAPL --max-symbols 2

Full universe extension (long run)::

    cd /Users/robzingale/trading_bot && PYTHONUNBUFFERED=1 .venv/bin/python \\
        RenTech/data_pipeline/download_alpaca_rth_minutes.py \\
        --start 2025-01-02 --end today
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
import warnings
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timedelta, timezone
from pathlib import Path

import pandas as pd

_REPO = Path(__file__).resolve().parents[2]
if str(_REPO) not in sys.path:
    sys.path.insert(0, str(_REPO))

from RenTech.strategy_stack.alpaca_minute_loader import DEFAULT_ALPACA_RTH_DIR, list_parquet_symbols

# Reuse RTH grid builder from repo root script.
import importlib.util

_spec = importlib.util.spec_from_file_location(
    "rth_grid", _REPO / "0_build_rth_minute_grid.py"
)
_rth = importlib.util.module_from_spec(_spec)
assert _spec.loader is not None
_spec.loader.exec_module(_rth)

NAIVE_TIMESTAMP_TZ = _rth.NAIVE_TIMESTAMP_TZ
_resample_to_full_grid = _rth._resample_to_full_grid

LOGS = _REPO / "RenTech" / "data" / "logs"
DEFAULT_ENV = _REPO / "trading_bot_live" / ".env"
DEFAULT_STATE = LOGS / "alpaca_rth_download_state.json"

# Alpaca market-data API (Basic): stay under ~200 req/min.
DEFAULT_WORKERS = 4
DEFAULT_CHUNK_DAYS = 14
DEFAULT_SLEEP_SEC = 0.25
ALPACA_BAR_LIMIT = 10_000


class _InvalidAlpacaSymbol(Exception):
    pass


def _is_invalid_symbol_error(exc: BaseException) -> bool:
    msg = str(exc).lower()
    return "invalid symbol" in msg


def _load_env(path: Path) -> None:
    if not path.is_file():
        raise FileNotFoundError(f"Missing Alpaca env file: {path}")
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


def _parse_end(s: str) -> datetime:
    s = s.strip().lower()
    if s in ("today", "now"):
        # Basic plan: no bars in latest ~15 min; use prior close day if early.
        return datetime.now(timezone.utc).replace(hour=23, minute=59, second=0, microsecond=0)
    return pd.Timestamp(s).tz_localize("UTC").to_pydatetime()


def _alpaca_client():
    from alpaca.data.historical import StockHistoricalDataClient

    key = os.environ.get("APCA_API_KEY_ID") or os.environ.get("ALPACA_API_KEY_ID")
    secret = os.environ.get("APCA_API_SECRET_KEY") or os.environ.get("ALPACA_API_SECRET_KEY")
    if not key or not secret:
        raise RuntimeError("Set APCA_API_KEY_ID and APCA_API_SECRET_KEY in trading_bot_live/.env")
    return StockHistoricalDataClient(key, secret)


def _bars_df(symbol: str, start: datetime, end: datetime, *, feed) -> pd.DataFrame:
    from alpaca.data.requests import StockBarsRequest
    from alpaca.data.timeframe import TimeFrame

    client = _alpaca_client()
    req = StockBarsRequest(
        symbol_or_symbols=symbol,
        timeframe=TimeFrame.Minute,
        start=start,
        end=end,
        feed=feed,
        limit=10_000,
    )
    resp = client.get_stock_bars(req)
    if resp is None or resp.df is None or resp.df.empty:
        return pd.DataFrame()
    df = resp.df
    if isinstance(df.index, pd.MultiIndex):
        if symbol not in df.index.get_level_values(0):
            return pd.DataFrame()
        df = df.xs(symbol, level=0)
    rows = []
    for ts, row in df.iterrows():
        t = pd.Timestamp(ts)
        if t.tzinfo is not None:
            t = t.tz_convert("UTC").tz_localize(None)
        rows.append(
            {
                "datetime": t,
                "open": float(row["open"]),
                "high": float(row["high"]),
                "low": float(row["low"]),
                "close": float(row["close"]),
                "volume": float(row["volume"]),
            }
        )
    if not rows:
        return pd.DataFrame()
    out = pd.DataFrame(rows).sort_values("datetime").drop_duplicates("datetime")
    out = out.set_index("datetime")
    for c in ("open", "high", "low", "close", "volume"):
        out[c] = pd.to_numeric(out[c], errors="coerce")
    return out


def _fetch_symbol_raw(
    symbol: str,
    start: datetime,
    end: datetime,
    *,
    chunk_days: int,
    feed,
    sleep_sec: float,
) -> pd.DataFrame:
    """Fetch all minute bars; Alpaca caps at 10k bars/request so cursor-advance."""
    chunks: list[pd.DataFrame] = []
    t0 = start if start.tzinfo else start.replace(tzinfo=timezone.utc)
    end_utc = end if end.tzinfo else end.replace(tzinfo=timezone.utc)
    while t0 < end_utc:
        t1 = min(t0 + timedelta(days=int(chunk_days)), end_utc)
        try:
            part = _bars_df(symbol, t0, t1, feed=feed)
            if part.empty:
                t0 = t1
            else:
                chunks.append(part)
                last = pd.Timestamp(part.index.max())
                if last.tzinfo is None:
                    last = last.tz_localize("UTC")
                # Hit page cap → continue from next minute, not next calendar chunk.
                if len(part) >= ALPACA_BAR_LIMIT - 5:
                    t0 = (last + pd.Timedelta(minutes=1)).to_pydatetime()
                else:
                    t0 = t1
        except Exception as exc:
            if _is_invalid_symbol_error(exc):
                raise _InvalidAlpacaSymbol(str(exc)) from exc
            warnings.warn(f"{symbol} {t0.date()}–{t1.date()}: {exc}", stacklevel=1)
            t0 = t1
        if sleep_sec > 0:
            time.sleep(sleep_sec)
    if not chunks:
        return pd.DataFrame()
    out = pd.concat(chunks).sort_index()
    return out[~out.index.duplicated(keep="last")]


def _read_existing(path: Path) -> pd.DataFrame:
    if not path.is_file():
        return pd.DataFrame()
    df = pd.read_parquet(path)
    df.columns = [str(c).lower() for c in df.columns]
    df["datetime"] = pd.to_datetime(df["datetime"])
    return df.sort_values("datetime")


def _merge_and_grid(existing: pd.DataFrame, new_raw: pd.DataFrame) -> pd.DataFrame:
    if new_raw.empty:
        return existing
    new_raw = new_raw.copy()
    new_raw.index.name = "datetime"
    grid = _resample_to_full_grid(new_raw, naive_tz=NAIVE_TIMESTAMP_TZ)
    if grid.empty:
        return existing
    grid.index.name = "datetime"
    grid = grid.reset_index()
    if existing.empty:
        return grid.sort_values("datetime")
    combined = (
        pd.concat([existing, grid], ignore_index=True)
        .drop_duplicates(subset=["datetime"], keep="last")
        .sort_values("datetime")
    )
    return combined


def _write_parquet(df: pd.DataFrame, path: Path) -> None:
    out = df.copy()
    for c in ("open", "high", "low", "close", "volume"):
        out[c] = out[c].astype("float32")
    if "synthetic" not in out.columns:
        out["synthetic"] = 0
    out["synthetic"] = out["synthetic"].fillna(0).astype("int8")
    tmp = path.with_suffix(".parquet.tmp")
    out.to_parquet(tmp, index=False)
    tmp.replace(path)


def _effective_start(symbol: str, path: Path, cli_start: datetime) -> datetime:
    existing = _read_existing(path)
    if existing.empty:
        return cli_start
    last = pd.Timestamp(existing["datetime"].max())
    if last.tzinfo is None:
        last = last.tz_localize("UTC")
    else:
        last = last.tz_convert("UTC")
    # Next minute after last stored bar.
    nxt = (last + pd.Timedelta(minutes=1)).to_pydatetime()
    return max(cli_start, nxt.replace(tzinfo=timezone.utc))


def _process_symbol(
    symbol: str,
    *,
    data_dir: Path,
    cli_start: datetime,
    end: datetime,
    chunk_days: int,
    feed,
    sleep_sec: float,
) -> tuple[str, str, int, str | None]:
    path = data_dir / f"{symbol}.parquet"
    start = _effective_start(symbol, path, cli_start)
    if start >= end:
        existing = _read_existing(path)
        n = len(existing)
        return symbol, "skip_up_to_date", n, None
    try:
        existing = _read_existing(path)
        new_raw = _fetch_symbol_raw(
            symbol, start, end, chunk_days=chunk_days, feed=feed, sleep_sec=sleep_sec
        )
        if new_raw.empty:
            existing = _read_existing(path)
            if not existing.empty:
                last = pd.Timestamp(existing["datetime"].max())
                if last.tzinfo is None:
                    last = last.tz_localize("UTC")
                end_ts = pd.Timestamp(end)
                if end_ts.tzinfo is None:
                    end_ts = end_ts.tz_localize("UTC")
                # Already extended through window (or within 2 sessions) — don't retry forever.
                if last >= end_ts - pd.Timedelta(days=3):
                    return symbol, "skip_up_to_date", len(existing), str(last)
            return symbol, "no_new_bars", len(existing), "empty Alpaca response"
        merged = _merge_and_grid(existing, new_raw)
        if merged.empty:
            return symbol, "empty", 0, "grid merge empty"
        _write_parquet(merged, path)
        added = len(merged) - len(existing)
        new_max = merged["datetime"].max()
        return symbol, "ok", int(added), str(new_max)
    except _InvalidAlpacaSymbol as exc:
        return symbol, "skip_invalid", 0, str(exc)
    except Exception as exc:
        return symbol, "error", 0, str(exc)


def _load_state(path: Path) -> set[str]:
    if not path.is_file():
        return set()
    try:
        data = json.loads(path.read_text())
        return set(data.get("done", []))
    except json.JSONDecodeError:
        return set()


def _save_state(path: Path, done: set[str], meta: dict) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {"done": sorted(done), **meta}
    path.write_text(json.dumps(payload, indent=2) + "\n")


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--data-dir", type=Path, default=DEFAULT_ALPACA_RTH_DIR)
    ap.add_argument("--env-file", type=Path, default=DEFAULT_ENV)
    ap.add_argument("--start", default="2025-01-02", help="Fetch from (UTC date); skips if parquet newer")
    ap.add_argument("--end", default="today", help="End date (UTC) or 'today'")
    ap.add_argument("--symbols", default="", help="Comma-separated tickers (default: all existing parquets)")
    ap.add_argument("--max-symbols", type=int, default=0, help="Limit tickers (0 = all)")
    ap.add_argument("--workers", type=int, default=DEFAULT_WORKERS)
    ap.add_argument("--chunk-days", type=int, default=DEFAULT_CHUNK_DAYS)
    ap.add_argument("--sleep-sec", type=float, default=DEFAULT_SLEEP_SEC, help="Pause between API chunks")
    ap.add_argument(
        "--feed",
        default="iex",
        choices=("iex", "sip"),
        help="Alpaca data feed (Basic plan = iex)",
    )
    ap.add_argument("--resume", action="store_true", help="Skip symbols marked done in state file")
    ap.add_argument("--state-file", type=Path, default=DEFAULT_STATE)
    args = ap.parse_args()

    _load_env(args.env_file.expanduser().resolve())
    from alpaca.data.enums import DataFeed

    feed = DataFeed.SIP if args.feed == "sip" else DataFeed.IEX

    data_dir = args.data_dir.expanduser().resolve()
    data_dir.mkdir(parents=True, exist_ok=True)

    cli_start = pd.Timestamp(args.start).tz_localize("UTC").to_pydatetime()
    end = _parse_end(args.end)

    if args.symbols.strip():
        symbols = [s.strip().upper() for s in args.symbols.split(",") if s.strip()]
    else:
        symbols = list_parquet_symbols(data_dir)
    if args.max_symbols and args.max_symbols > 0:
        symbols = symbols[: int(args.max_symbols)]

    done = _load_state(args.state_file) if args.resume else set()
    todo = [s for s in symbols if s not in done]
    print(
        f"Alpaca RTH download — symbols={len(todo)}/{len(symbols)} "
        f"window={cli_start.date()}→{end.date()} feed={args.feed}",
        flush=True,
    )

    results: list[tuple[str, str, int, str | None]] = []
    t0 = time.time()
    workers = max(1, int(args.workers))

    with ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {
            ex.submit(
                _process_symbol,
                sym,
                data_dir=data_dir,
                cli_start=cli_start,
                end=end,
                chunk_days=int(args.chunk_days),
                feed=feed,
                sleep_sec=float(args.sleep_sec),
            ): sym
            for sym in todo
        }
        for i, fut in enumerate(as_completed(futs), 1):
            sym = futs[fut]
            try:
                row = fut.result()
            except Exception as exc:
                row = (sym, "error", 0, str(exc))
            results.append(row)
            status, added, detail = row[1], row[2], row[3]
            if status == "ok":
                done.add(sym)
                if i % 25 == 0 or sym in ("SPY", "AAPL"):
                    print(f"  [{i}/{len(todo)}] {sym} +{added} rows → max {detail}", flush=True)
            elif status in ("skip_up_to_date", "skip_invalid", "no_new_bars"):
                done.add(sym)
                if sym in ("SPY",) and status == "skip_up_to_date":
                    print(f"  [{i}/{len(todo)}] {sym} up to date ({added} rows)", flush=True)
            if i % 100 == 0:
                _save_state(
                    args.state_file,
                    done,
                    {
                        "last_update_utc": datetime.now(timezone.utc).isoformat(),
                        "processed": i,
                        "total": len(todo),
                    },
                )

    _save_state(
        args.state_file,
        done,
        {
            "last_update_utc": datetime.now(timezone.utc).isoformat(),
            "processed": len(results),
            "total": len(todo),
            "elapsed_sec": round(time.time() - t0, 1),
        },
    )

    ok = sum(1 for r in results if r[1] == "ok")
    skip = sum(1 for r in results if r[1] == "skip_up_to_date")
    empty = sum(1 for r in results if r[1] in ("no_new_bars", "empty"))
    err = sum(1 for r in results if r[1] == "error")
    added_total = sum(r[2] for r in results if r[1] == "ok")

    print(
        f"\nDone in {time.time()-t0:.0f}s — ok={ok} skip={skip} empty={empty} err={err} "
        f"rows_added={added_total}",
        flush=True,
    )
    if err:
        print("Errors (first 10):", flush=True)
        for sym, st, _, det in [r for r in results if r[1] == "error"][:10]:
            print(f"  {sym}: {det}", flush=True)

    # Verify SPY max date
    spy_path = data_dir / "SPY.parquet"
    if spy_path.is_file():
        spy = _read_existing(spy_path)
        print(f"SPY max datetime: {spy['datetime'].max()}  rows={len(spy)}", flush=True)


if __name__ == "__main__":
    main()