#!/usr/bin/env python3
"""Massive reference data for every stock that appears in data/massive/selection (operator 2026-10-01: the plan ends
2026-11-01 — take everything useful while we have it). One file per stock and kind under data/massive/ref/<kind>/,
resumable (files on disk are skipped).

  details        ticker details as of the stock's last selected day (SIC code + description, address, market cap,
                 shares outstanding, list date, ...) — segments (Chinese stocks, oil, biotech) and size
  float          current free float
  events         ticker changes (renamed stocks)
  news           every news item 2021-10-01 → now with its tickers and per-ticker sentiment ("insights")
  short_interest bi-monthly short interest
  short_volume   daily short volume
  financials     quarterly financials (share counts → dilution)
  ipos           every IPO since 2021-10-01 (one file)

    python scripts/massive_reference.py [kinds...] [--workers 5]
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
import urllib.parse
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from massive_download import get, write_gz, ROOT, WINDOW_START  # noqa: E402

REF = ROOT / "ref"
START = WINDOW_START


def selection_tickers() -> dict[str, str]:
    """ticker → its last selected day."""
    last = {}
    for p in sorted((ROOT / "selection").glob("*.json")):
        for s in json.load(open(p))["symbols"]:
            last[s] = p.name[:10]
    return last


def paged(url: str, max_pages: int = 2000) -> list:
    rows, pages = [], 0
    while url and pages < max_pages:
        js = get(url)
        if js.get("http") == 403:
            break
        res = js.get("results")
        if isinstance(res, list):
            rows += res
        elif res:
            rows.append(res)
        url = js.get("next_url")
        pages += 1
    return rows


def fetch(kind: str, sym: str, last_day: str):
    q = urllib.parse.quote(sym, safe="")
    if kind == "details":
        return get(f"/v3/reference/tickers/{q}?date={last_day}").get("results") or {}
    if kind == "float":
        return paged(f"/stocks/vX/float?ticker={q}&limit=100")
    if kind == "events":
        return get(f"/vX/reference/tickers/{q}/events").get("results") or {}
    if kind == "news":
        return [{k: r.get(k) for k in ("id", "published_utc", "title", "description", "tickers", "keywords", "insights", "article_url")}
                | {"publisher": (r.get("publisher") or {}).get("name")}
                for r in paged(f"/v2/reference/news?ticker={q}&published_utc.gte={START}&limit=1000&order=asc&sort=published_utc")]
    if kind == "short_interest":
        return paged(f"/stocks/v1/short-interest?ticker={q}&limit=50000")
    if kind == "short_volume":
        return paged(f"/stocks/v1/short-volume?ticker={q}&limit=50000")
    if kind == "financials":
        return paged(f"/vX/reference/financials?ticker={q}&timeframe=quarterly&limit=100&filing_date.gte={START}")
    raise ValueError(kind)


def run(kind: str, workers: int) -> None:
    tick = selection_tickers()
    todo = [(s, d) for s, d in sorted(tick.items()) if not (REF / kind / f"{s}.json.gz").exists()]
    print(f"{kind}: {len(tick):,} stocks, {len(todo):,} to fetch", flush=True)
    t0, done, errors = time.time(), 0, 0

    def job(s, d):
        write_gz(REF / kind / f"{s}.json.gz", fetch(kind, s, d))

    with ThreadPoolExecutor(workers) as pool:
        futs = [pool.submit(job, s, d) for s, d in todo]
        for f in as_completed(futs):
            try:
                f.result()
            except Exception as e:  # noqa: BLE001 — retried on the next run
                errors += 1
                if errors <= 5:
                    print(f"  {kind} error: {type(e).__name__}: {str(e)[:120]}", flush=True)
            done += 1
            if done % 1000 == 0:
                print(f"  {kind} {done:,}/{len(todo):,} ({(time.time() - t0) / 60:.0f} min)", flush=True)
    print(f"{kind} done: {done:,}, errors {errors}", flush=True)


def ipos() -> None:
    rows = paged(f"/vX/reference/ipos?listing_date.gte={START}&limit=1000&order=asc&sort=listing_date")
    write_gz(REF / "ipos.json.gz", rows)
    print(f"ipos: {len(rows):,}")


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("kinds", nargs="*", default=["ipos", "details", "events", "float", "financials", "short_interest", "news", "short_volume"])
    ap.add_argument("--workers", type=int, default=5)
    a = ap.parse_args()
    for k in a.kinds:
        ipos() if k == "ipos" else run(k, a.workers)