"""Fetch FINRA REG SHO daily short-sale volume for the research universes.

One public file per trading day covers every NMS symbol (no auth):

    https://cdn.finra.org/equity/regsho/daily/CNMSshvolYYYYMMDD.txt
    Date|Symbol|ShortVolume|ShortExemptVolume|TotalVolume|Market

Writes data/research/short_volume.pkl: {"short": DataFrame, "total": DataFrame}
(float32, date x symbol), filtered to the symbols of data/us.csv plus the US
mid/small panels (us_mid / us_midsmall). Resume-safe: re-running fetches only
days that are missing from the store. Network goes through QUANT_MARKET_PROXY
when set (same convention as the other fetchers).

    PYTHONPATH=. uv run python -m research.lab.fetch_shortvolume [--days-limit N]
"""
from __future__ import annotations

import json
import os
import pickle
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

import pandas as pd
import requests

from research.lab.fetch_us import _proxy_env

OUT = os.path.join("data", "research", "short_volume.pkl")
SYMBOLS_CACHE = os.path.join("data", "research", "shortvolume_symbols.json")
URL = "https://cdn.finra.org/equity/regsho/daily/CNMSshvol{day}.txt"
HEADERS = {"User-Agent": "Mozilla/5.0 (research; contact local)"}
WORKERS = 6


def universe_symbols() -> list[str]:
    if os.path.exists(SYMBOLS_CACHE):
        return json.load(open(SYMBOLS_CACHE))
    syms = set()
    us = pd.read_csv("data/us.csv", index_col=0, nrows=0)
    syms.update(us.columns)
    for path in ("data/research/us_mid.panel.pkl", "data/research/us_midsmall.panel.pkl"):
        if os.path.exists(path):
            with open(path, "rb") as fh:
                panel = pickle.load(fh)
            cols = getattr(panel, "close").columns if hasattr(panel, "close") else None
            if cols is None and isinstance(panel, dict):
                cols = panel["close"].columns
            syms.update(str(c).replace(".", "-") for c in cols)
    syms.discard("")
    out = sorted(syms)
    json.dump(out, open(SYMBOLS_CACHE, "w"))
    return out


def trading_days() -> list[str]:
    close = pd.read_csv("data/us.csv", index_col=0, usecols=[0], parse_dates=True)
    days = close.index
    if os.path.exists("data/etfs.csv"):
        etfs = pd.read_csv("data/etfs.csv", index_col=0, usecols=[0], parse_dates=True)
        days = days.union(etfs.index)
    return [d.strftime("%Y%m%d") for d in days]


def fetch_day(day: str, symbols: set[str]) -> tuple[str, dict, dict] | None:
    url = URL.format(day=day)
    for attempt in range(4):
        try:
            r = requests.get(url, headers=HEADERS, timeout=40)
            if r.status_code in (403, 404):
                # FINRA's public archive keeps only ~8 years; older days are gone.
                return day, {}, {}
            r.raise_for_status()
            short, total = {}, {}
            for line in r.text.splitlines()[1:]:
                parts = line.split("|")
                if len(parts) < 5:
                    continue
                sym = parts[1]
                if sym not in symbols:
                    continue
                try:
                    s, t = float(parts[2]), float(parts[4])
                except ValueError:
                    continue
                short[sym] = short.get(sym, 0.0) + s
                total[sym] = total.get(sym, 0.0) + t
            return day, short, total
        except Exception as exc:  # noqa: BLE001
            if attempt == 3:
                print(f"  FAILED {day}: {exc}")
            else:
                import time
                time.sleep(1.5 * (attempt + 1))
    return None


def main() -> None:
    _proxy_env()
    limit = None
    if "--days-limit" in sys.argv:
        limit = int(sys.argv[sys.argv.index("--days-limit") + 1])
    symbols = set(universe_symbols())
    print(f"universe symbols: {len(symbols)}")
    days = trading_days()

    prev_short = prev_total = None
    have: set[str] = set()
    if os.path.exists(OUT):
        with open(OUT, "rb") as fh:
            d = pickle.load(fh)
        prev_short, prev_total = d["short"], d["total"]
        have = {pd.Timestamp(i).strftime("%Y%m%d") for i in prev_short.index}
    todo = [d for d in days if d not in have]
    if limit:
        todo = todo[:limit]
    print(f"days: {len(days)}, already stored: {len(have)}, fetching: {len(todo)}")

    short_rows: dict[str, dict] = {}
    total_rows: dict[str, dict] = {}
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = {ex.submit(fetch_day, d, symbols): d for d in todo}
        done = 0
        for fut in as_completed(futs):
            res = fut.result()
            done += 1
            if res is not None:
                day, short, total = res
                short_rows[day] = short
                total_rows[day] = total
            if done % 200 == 0:
                print(f"  {done}/{len(todo)} fetched")

    new_short = pd.DataFrame.from_dict(short_rows, orient="index")
    new_total = pd.DataFrame.from_dict(total_rows, orient="index")
    for df in (new_short, new_total):
        df.index = pd.to_datetime(df.index, format="%Y%m%d")
    frames_short = [f for f in (prev_short, new_short) if f is not None and len(f)]
    frames_total = [f for f in (prev_total, new_total) if f is not None and len(f)]
    df_short = pd.concat(frames_short).sort_index()
    df_total = pd.concat(frames_total).sort_index()
    with open(OUT, "wb") as fh:
        pickle.dump({"short": df_short.astype("float32"), "total": df_total.astype("float32")}, fh)
    print(f"stored {df_short.shape} short / {df_total.shape} total -> {OUT}")


if __name__ == "__main__":
    main()