#!/usr/bin/env python3
"""Massive (ex-Polygon) 5-year US stock history — the raw data behind the balance-method and lanes tests.
Your own Massive Stocks plan key is read from the MASSIVE_API_KEY environment variable and sent only in the
Authorization header. Files go under MASSIVE_ROOT (default ./massive_data). Run the phases in this order:
  daily, reference, splits, select, minute, pmcand, rthdaily   (python massive_download.py <phase>; 'status' shows progress)
The original notes follow.

Massive (ex-Polygon) Stocks Starter history → data/massive/ (operator 2026-10-01: $29 plan, MASSIVE_API_KEY in
.env.local). The plan is a ROLLING 5 years (first allowed day 2021-10-02 on 2026-10-01): every phase works oldest
day first, because the oldest days expire one per day.

Phases (each resumable; files already on disk are skipped):
  daily      grouped daily bars for the whole US stock market, one call per trading day → grouped/<day>.json.gz
             (regular-session open/high/low/close = the official open and close; volume = the consolidated total)
  reference  every stock ticker, active and delisted, with its type (CS, ADRC, ETF, WARRANT, ...) → reference/
  splits     every split since the window start → reference/splits.json.gz
  select     the day's movers, same rules as the IBKR store (scripts/market_history_download.py select_day) but on the
             regular-session daily bar (Massive's daily high/low exclude premarket and after-hours, so a stock that
             spiked only in premarket is missed): high >= +10% over the prior close, or range >= 20%, or close moved
             >= 10%, with >= $250k traded and a high >= $0.50; common stocks and ADRs only (reference types CS / ADRC);
             the prior close is split-adjusted; + the day after each mover; + SPY / QQQ / IWM → selection/<day>.json
  minute     1-minute bars 04:00-20:00 (as traded) for every selected stock-day AND its previous trading day (the
             Flush lane needs yesterday's zones), one call per run of consecutive days → minute/<day>/<SYM>.json.gz in
             the IBKR store's format (timestamp, open, high, low, close, volume), so scripts/market_store.py can read it
  rthdaily   per-stock regular-session daily bars from the grouped files (selected stocks only) → daily/rth/<SYM>.json.gz
  pmcand     completes the selection with stocks that moved only outside regular hours (operator 2026-10-01: no selection
             bias — every stock that could have qualified in real time must be in the list). Candidates = common stocks /
             ADRs at $0.50-$50 not yet selected that opened 5%+ above the prior close with $250k+ traded, or traded 3x
             their 20-day median dollars (and $500k+); their minutes (that day + the day before) are fetched and a stock
             whose high since 16:00 the day before is 10%+ over the prior close (or whose 04:00-20:00 range is 20%+), with
             $250k+ traded, joins the day's selection ("eth_high>=+10%") and the next day's ("day_after_mover")
  status     what is on disk

The key is read from .env.local and sent in the Authorization header (never in a URL or a log line).

    python scripts/massive_download.py daily [--from 2021-10-04 --to 2026-09-30 --workers 8]
    python scripts/massive_download.py reference
    python scripts/massive_download.py splits
    python scripts/massive_download.py select
    python scripts/massive_download.py minute [--workers 8]
    python scripts/massive_download.py rthdaily
    python scripts/massive_download.py pmcand [--from --to --workers 10]
    python scripts/massive_download.py status
"""
from __future__ import annotations

import argparse
import gzip
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from collections import defaultdict
from datetime import datetime, timezone
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import date, timedelta
from pathlib import Path

ROOT = Path(os.environ.get("MASSIVE_ROOT", "massive_data"))
API = "https://api.massive.com"
# the plan is a rolling 5 years: start a few days inside it (override with MASSIVE_FROM=YYYY-MM-DD)
WINDOW_START = os.environ.get("MASSIVE_FROM", (date.today() - timedelta(days=5 * 365 - 5)).isoformat())


def _key() -> str:
    key = os.environ.get("MASSIVE_API_KEY", "").strip()
    if key:
        return key
    if os.path.exists(".env.local"):
        for line in open(".env.local"):
            if line.startswith("MASSIVE_API_KEY="):
                return line.split("=", 1)[1].strip().strip('"').strip("'")
    raise SystemExit("set MASSIVE_API_KEY (export MASSIVE_API_KEY=...) — your own Massive Stocks plan key")


KEY = _key()


def get(path_or_url: str, tries: int = 6) -> dict:
    url = path_or_url if path_or_url.startswith("http") else API + path_or_url
    delay = 2.0
    for i in range(tries):
        req = urllib.request.Request(url, headers={"Authorization": f"Bearer {KEY}"})
        try:
            with urllib.request.urlopen(req, timeout=90) as r:
                return json.load(r)
        except urllib.error.HTTPError as e:
            if e.code == 403:
                body = e.read()[:200].decode(errors="replace")
                return {"status": "NOT_AUTHORIZED", "http": 403, "message": body}
            if e.code in (429, 500, 502, 503, 504) and i < tries - 1:
                time.sleep(delay)
                delay *= 2
                continue
            raise
        except (urllib.error.URLError, TimeoutError, ConnectionError):
            if i < tries - 1:
                time.sleep(delay)
                delay *= 2
                continue
            raise
    return {}


def write_gz(path: Path, obj) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_suffix(path.suffix + ".tmp")
    with gzip.open(tmp, "wt") as fh:
        json.dump(obj, fh, separators=(",", ":"))
    tmp.replace(path)


def weekdays(frm: str, to: str) -> list[str]:
    d, end, out = date.fromisoformat(frm), date.fromisoformat(to), []
    while d <= end:
        if d.weekday() < 5:
            out.append(d.isoformat())
        d += timedelta(days=1)
    return out


def one_daily(day: str) -> tuple[str, int | str]:
    path = ROOT / "grouped" / f"{day}.json.gz"
    if path.exists():
        return day, "skip"
    js = get(f"/v2/aggs/grouped/locale/us/market/stocks/{day}?adjusted=false")
    if js.get("http") == 403:
        return day, "expired"
    rows = js.get("results") or []
    write_gz(path, rows)           # an empty list = market holiday (kept so the day is not asked again)
    return day, len(rows)


def daily(frm: str, to: str, workers: int) -> None:
    days = weekdays(max(frm, WINDOW_START), to)
    done = n_rows = 0
    t0 = time.time()
    with ThreadPoolExecutor(workers) as pool:
        futs = {pool.submit(one_daily, d): d for d in days}         # submitted oldest first
        for f in as_completed(futs):
            day, res = f.result()
            done += 1
            if isinstance(res, int):
                n_rows += res
            if res == "expired":
                print(f"  {day}: outside the plan's window", flush=True)
            if done % 100 == 0:
                print(f"  daily {done}/{len(days)} ({time.time() - t0:.0f}s)", flush=True)
    print(f"daily done: {done} weekdays, {n_rows:,} rows written this run")


def reference() -> None:
    for active in ("true", "false"):
        path = ROOT / "reference" / f"tickers_{'active' if active == 'true' else 'inactive'}.json.gz"
        if path.exists():
            print(f"  {path.name} exists")
            continue
        rows, url, pages = [], f"/v3/reference/tickers?market=stocks&active={active}&limit=1000&order=asc&sort=ticker", 0
        while url:
            js = get(url)
            rows += [{k: r.get(k) for k in ("ticker", "type", "primary_exchange", "name", "active", "delisted_utc", "cik",
                                            "composite_figi", "share_class_figi", "list_date")} for r in js.get("results") or []]
            url = js.get("next_url")
            pages += 1
            if pages % 20 == 0:
                print(f"  tickers active={active}: {len(rows):,} after {pages} pages", flush=True)
        write_gz(path, rows)
        print(f"  {path.name}: {len(rows):,} tickers")


def splits() -> None:
    path = ROOT / "reference" / "splits.json.gz"
    rows, url = [], f"/v3/reference/splits?execution_date.gte={WINDOW_START}&limit=1000&order=asc&sort=execution_date"
    while url:
        js = get(url)
        rows += js.get("results") or []
        url = js.get("next_url")
    write_gz(path, rows)
    print(f"  splits: {len(rows):,}")


BENCHMARKS = ("SPY", "QQQ", "IWM")
KEEP_TYPES = {"CS", "ADRC"}


@__import__("functools").lru_cache(maxsize=1)
def _trading_days() -> tuple:
    out = []
    for p in sorted((ROOT / "grouped").glob("*.json.gz")):
        if len(json.load(gzip.open(p, "rt"))) > 1000:   # a market holiday has an empty (or near-empty) file
            out.append(p.name[:10])
    return tuple(out)


def trading_days() -> list[str]:
    return list(_trading_days())


def _types() -> dict[str, set]:
    types = defaultdict(set)
    for name in ("tickers_active", "tickers_inactive"):
        p = ROOT / "reference" / f"{name}.json.gz"
        if p.exists():
            for r in json.load(gzip.open(p, "rt")):
                types[r["ticker"]].add(r.get("type"))
    return types


def _split_factor() -> dict[tuple, float]:
    """(ticker, execution day) → multiply the previous close by this to compare it with the split day's prices."""
    p = ROOT / "reference" / "splits.json.gz"
    out = {}
    for r in (json.load(gzip.open(p, "rt")) if p.exists() else []):
        if r.get("split_to") and r.get("split_from"):
            out[(r["ticker"], r["execution_date"])] = float(r["split_from"]) / float(r["split_to"])
    return out


def _looks_common(t: str) -> bool:
    return "." not in t and not (len(t) == 5 and t[-1] in "WURZ")


def select() -> None:
    days = trading_days()
    types, splits_f = _types(), _split_factor()
    if not types:
        raise SystemExit("reference tickers missing: run the reference phase first")
    prev_rows, prev_sel = None, {}
    n_core = 0
    for day in days:
        rows = {r["T"]: r for r in json.load(gzip.open(ROOT / "grouped" / f"{day}.json.gz", "rt"))}
        out = {s: ["benchmark"] for s in BENCHMARKS}
        if prev_rows is not None:
            for sym, r in rows.items():
                ts = types.get(sym)
                if (ts and not (ts & KEEP_TYPES)) or (not ts and not _looks_common(sym)):
                    continue
                p = prev_rows.get(sym)
                if not p or not p.get("c"):
                    continue
                prior = float(p["c"]) * splits_f.get((sym, day), 1.0)
                hi, lo, close, vol = float(r["h"]), float(r["l"]), float(r["c"]), float(r.get("v") or 0)
                if prior <= 0 or hi < 0.5 or vol * close < 250_000:
                    continue
                why = []
                if hi / prior >= 1.10:
                    why.append("high>=+10%_vs_prior_close")
                if lo > 0 and hi / lo >= 1.20:
                    why.append("range>=20%")
                if abs(close / prior - 1.0) >= 0.10:
                    why.append("close_moved>=10%")
                if why:
                    out.setdefault(sym, []).extend(why)
                    n_core += 1
            for sym, why in prev_sel.items():
                if "benchmark" not in why and "day_after_mover" not in why and sym in rows:
                    out.setdefault(sym, []).append("day_after_mover")
        path = ROOT / "selection" / f"{day}.json"
        path.parent.mkdir(parents=True, exist_ok=True)
        json.dump({"day": day, "count": len(out), "symbols": out}, open(path, "w"))
        prev_rows, prev_sel = rows, out
    print(f"select: {len(days)} days, {n_core:,} mover-days (+ day-after names)")


def _runs(day_idx: list[int]) -> list[tuple[int, int]]:
    runs, start, last = [], None, None
    for i in sorted(set(day_idx)):
        if start is None:
            start = last = i
        elif i == last + 1 and i - start < 40:
            last = i
        else:
            runs.append((start, last))
            start = last = i
    if start is not None:
        runs.append((start, last))
    return runs


def _et_day(ms_ts: int) -> str:
    return (datetime.fromtimestamp(ms_ts / 1000, timezone.utc) - timedelta(hours=4)).date().isoformat()


def one_minute_run(sym: str, d0: str, d1: str, want: list[str]) -> int:
    q = urllib.parse.quote(sym, safe="")
    url = f"/v2/aggs/ticker/{q}/range/1/minute/{d0}/{d1}?adjusted=false&sort=asc&limit=50000"
    bars = []
    while url:
        js = get(url)
        if js.get("http") == 403:
            return -1
        bars += js.get("results") or []
        url = js.get("next_url")
    by_day = defaultdict(list)
    for b in bars:
        ts = datetime.fromtimestamp(b["t"] / 1000, timezone.utc)
        # ET date: EDT/EST handled by checking both offsets against the session window is overkill — bars are 04:00-20:00
        # ET, i.e. 08:00-01:00 UTC; a bar's ET date is its UTC date shifted back 4-5 h, never across midnight before 04:00
        et_day = (ts - timedelta(hours=5)).date().isoformat() if ts.hour < 5 else ts.date().isoformat()
        by_day[et_day].append({"timestamp": ts.strftime("%Y-%m-%dT%H:%M:%SZ"), "open": b["o"], "high": b["h"], "low": b["l"],
                               "close": b["c"], "volume": b["v"]})
    for d in want:
        write_gz(ROOT / "minute" / d / f"{sym}.json.gz", by_day.get(d, []))
    return len(bars)


def minute(workers: int) -> None:
    days = trading_days()
    idx = {d: i for i, d in enumerate(days)}
    want = defaultdict(set)
    for d in days:
        p = ROOT / "selection" / f"{d}.json"
        if not p.exists():
            continue
        for sym in json.load(open(p))["symbols"]:
            i = idx[d]
            want[sym].add(i)
            if i > 0:
                want[sym].add(i - 1)                    # the previous session: yesterday's zones / prior RTH close
    jobs = []
    for sym, ids in want.items():
        for a, b in _runs(list(ids)):
            dl = days[a:b + 1]
            if all((ROOT / "minute" / d / f"{sym}.json.gz").exists() for d in dl):
                continue
            jobs.append((days[a], sym, days[b], dl))
    jobs.sort()                                          # oldest first (the rolling window expires at the old end)
    print(f"minute: {sum(len(v) for v in want.values()):,} stock-days wanted, {len(jobs):,} calls to make", flush=True)
    t0, done, nbars, expired = time.time(), 0, 0, 0
    with ThreadPoolExecutor(workers) as pool:
        futs = [pool.submit(one_minute_run, sym, d0, d1, dl) for d0, sym, d1, dl in jobs]
        for f in as_completed(futs):
            try:
                n = f.result()
            except Exception as e:  # noqa: BLE001 — logged, retried on the next run
                print(f"  error: {type(e).__name__}", flush=True)
                continue
            done += 1
            if n < 0:
                expired += 1
            else:
                nbars += n
            if done % 2000 == 0:
                el = time.time() - t0
                print(f"  minute {done:,}/{len(jobs):,} calls, {nbars:,} bars, {el / 60:.0f} min, ~{(len(jobs) - done) * el / done / 60:.0f} min left", flush=True)
    print(f"minute done: {done:,} calls, {nbars:,} bars, {expired} outside the window")


def rthdaily() -> None:
    syms = set()
    for p in (ROOT / "selection").glob("*.json"):
        syms |= set(json.load(open(p))["symbols"])
    series = defaultdict(list)
    for day in trading_days():
        for r in json.load(gzip.open(ROOT / "grouped" / f"{day}.json.gz", "rt")):
            if r["T"] in syms:
                series[r["T"]].append({"timestamp": f"{day}T00:00:00Z", "open": r["o"], "high": r["h"], "low": r["l"],
                                       "close": r["c"], "volume": r.get("v")})
    for sym, rows in series.items():
        write_gz(ROOT / "daily" / "rth" / f"{sym}.json.gz", rows)
    print(f"rthdaily: {len(series):,} stocks")


def pm_candidates(frm: str, to: str) -> dict[str, list[str]]:
    import statistics
    from collections import deque
    days = trading_days()
    types = _types()
    hist = defaultdict(lambda: deque(maxlen=20))
    prev, out = None, {}
    for d in days:
        rows = {r["T"]: r for r in json.load(gzip.open(ROOT / "grouped" / f"{d}.json.gz", "rt"))}
        if prev is not None and frm <= d <= to:
            sel = set(json.load(open(ROOT / "selection" / f"{d}.json"))["symbols"])
            cands = []
            for sym, r in rows.items():
                ts = types.get(sym)
                if sym in sel or (ts and not (ts & KEEP_TYPES)) or (not ts and not _looks_common(sym)):
                    continue
                p = prev.get(sym)
                if not p or not p.get("c") or not (0.5 <= float(p["c"]) <= 50):
                    continue
                usd = float(r["c"]) * float(r.get("v") or 0)
                h = hist[sym]
                med = statistics.median(h) if len(h) >= 5 else None
                gap = float(r["o"]) / float(p["c"]) - 1
                if (gap >= 0.05 and usd >= 250_000) or (med and usd >= 3 * med and usd >= 500_000):
                    cands.append(sym)
            out[d] = cands
        for sym, r in rows.items():
            hist[sym].append(float(r["c"]) * float(r.get("v") or 0))
        prev = rows
    return out


def _eth_qualifies(sym: str, day: str, prev_day: str, prior: float) -> bool:
    from packages.mission_engine import lane_predicates as lp
    def bars(d):
        p = ROOT / "minute" / d / f"{sym}.json.gz"
        return lp.normalize_bars(json.load(gzip.open(p, "rt")), d) if p.exists() else []
    y = [b for b in bars(prev_day) if b["hm"] >= "16:00"]
    t = bars(day)
    if not t or not prior:
        return False
    usd = sum(b["close"] * b["volume"] for b in t)
    hi = max(b["high"] for b in y + t)
    lo = min(b["low"] for b in t)
    return usd >= 250_000 and max(b["high"] for b in t) >= 0.5 and (hi / prior >= 1.10 or (lo > 0 and max(b["high"] for b in t) / lo >= 1.20))


def pmcand(frm: str, to: str, workers: int) -> None:
    days = trading_days()
    idx = {d: i for i, d in enumerate(days)}
    cands = pm_candidates(frm, to)
    print(f"pmcand {frm}..{to}: {sum(len(v) for v in cands.values()):,} candidate stock-days", flush=True)
    want = defaultdict(set)
    for d, syms in cands.items():
        for sym in syms:
            want[sym] |= {idx[d], idx[d] - 1}
    jobs = []
    for sym, ids in want.items():
        for a, b in _runs(list(ids)):
            dl = days[a:b + 1]
            if not all((ROOT / "minute" / x / f"{sym}.json.gz").exists() for x in dl):
                jobs.append((days[a], sym, days[b], dl))
    jobs.sort()
    print(f"  {len(jobs):,} calls", flush=True)
    t0, done = time.time(), 0
    with ThreadPoolExecutor(workers) as pool:
        for f in as_completed([pool.submit(one_minute_run, sym, d0, d1, dl) for d0, sym, d1, dl in jobs]):
            try:
                f.result()
            except Exception as e:  # noqa: BLE001
                print(f"  error: {type(e).__name__}", flush=True)
            done += 1
            if done % 5000 == 0:
                el = time.time() - t0
                print(f"  pmcand {done:,}/{len(jobs):,} calls, {el / 60:.0f} min, ~{(len(jobs) - done) * el / done / 60:.0f} min left", flush=True)
    splits_f = _split_factor()
    added, after = defaultdict(list), defaultdict(list)
    for d, syms in cands.items():
        pd = days[idx[d] - 1]
        prev_rows = {r["T"]: r for r in json.load(gzip.open(ROOT / "grouped" / f"{pd}.json.gz", "rt"))}
        for sym in syms:
            p = prev_rows.get(sym)
            prior = float(p["c"]) * splits_f.get((sym, d), 1.0) if p else None
            if _eth_qualifies(sym, d, pd, prior):
                added[d].append(sym)
                if idx[d] + 1 < len(days):
                    after[days[idx[d] + 1]].append(sym)
    # the day after each new mover needs its minutes too
    extra = []
    for d, syms in after.items():
        for sym in syms:
            if not (ROOT / "minute" / d / f"{sym}.json.gz").exists():
                extra.append((d, sym, d, [d]))
    with ThreadPoolExecutor(workers) as pool:
        list(pool.map(lambda j: one_minute_run(j[1], j[0], j[2], j[3]), extra))
    n_add = 0
    for d in sorted(set(added) | set(after)):
        p = ROOT / "selection" / f"{d}.json"
        js = json.load(open(p))
        for sym in added.get(d, []):
            js["symbols"].setdefault(sym, []).append("eth_high>=+10%")
            n_add += 1
        for sym in after.get(d, []):
            w = js["symbols"].setdefault(sym, [])
            if "day_after_mover" not in w:
                w.append("day_after_mover")
        js["count"] = len(js["symbols"])
        json.dump(js, open(p, "w"))
    print(f"pmcand done: {n_add:,} stock-days added as outside-hours movers (+ their day after)")


def status() -> None:
    g = sorted((ROOT / "grouped").glob("*.json.gz")) if (ROOT / "grouped").exists() else []
    print(f"grouped daily files: {len(g)}" + (f" ({g[0].name[:10]} .. {g[-1].name[:10]})" if g else ""))
    for p in sorted((ROOT / "reference").glob("*.json.gz")) if (ROOT / "reference").exists() else []:
        print(f"  {p.name}: {len(json.load(gzip.open(p, 'rt'))):,}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("phase", choices=["daily", "reference", "splits", "select", "minute", "rthdaily", "pmcand", "status"])
    ap.add_argument("--from", dest="frm", default=WINDOW_START)
    ap.add_argument("--to", default=(date.today() - timedelta(days=1)).isoformat())
    ap.add_argument("--workers", type=int, default=8)
    a = ap.parse_args()
    {"daily": lambda: daily(a.frm, a.to, a.workers), "reference": reference, "splits": splits, "select": select,
     "minute": lambda: minute(a.workers), "rthdaily": rthdaily, "pmcand": lambda: pmcand(a.frm, a.to, a.workers), "status": status}[a.phase]()