"""Fetch SEC fails-to-deliver (FTD) data for the research universes.

Two public ZIP files per month (a = 1st half, b = 2nd half):

    https://www.sec.gov/files/data/fails-deliver-data/cnsfailsYYYYMM[a|b].zip

Each archive holds a pipe-delimited txt with one row per symbol and
settlement date: SETTLEMENT DATE|CUSIP|SYMBOL|QUANTITY (FAILS)|DESCRIPTION|
PRICE|. The archive covers 2018-01 onward on the current URL scheme (older
months 404; SEC also lags ~3 weeks for the newest half-month).

Writes data/research/fails.pkl: {"shares": DataFrame, "value": DataFrame}
(float32, settlement-date x symbol), filtered to the symbols of data/us.csv
plus the US mid/small panels. Resume-safe; UA follows the repo convention.

    PYTHONPATH=. uv run python -m research.lab.fetch_fails
"""
from __future__ import annotations

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

import pandas as pd
import requests

from research.lab.fetch_us import _proxy_env
from research.lab.fetch_shortvolume import SYMBOLS_CACHE, universe_symbols

OUT = os.path.join("data", "research", "fails.pkl")
URL = "https://www.sec.gov/files/data/fails-deliver-data/cnsfails{ym}{half}.zip"
UA = {"User-Agent": "quant-research/0.1 gahow@example.com", "Accept-Encoding": "gzip"}
WORKERS = 4


def candidates() -> list[str]:
    out = []
    for y in range(2018, 2027):
        for m in range(1, 13):
            for half in ("a", "b"):
                out.append(f"{y}{m:02d}{half}")
    return out


def fetch_one(ym: str, symbols: set[str]) -> tuple[str, dict, dict] | None:
    url = URL.format(ym=ym[:6], half=ym[6])
    for attempt in range(4):
        try:
            r = requests.get(url, headers=UA, timeout=60)
            if r.status_code == 404:
                return ym, {}, {}
            r.raise_for_status()
            zf = zipfile.ZipFile(io.BytesIO(r.content))
            shares: dict = {}
            value: dict = {}
            for name in zf.namelist():
                with zf.open(name) as fh:
                    text = io.TextIOWrapper(fh, encoding="latin-1")
                    header = True
                    for line in text:
                        if header:
                            header = False
                            continue
                        parts = line.rstrip("\n").split("|")
                        if len(parts) < 6:
                            continue
                        date, sym = parts[0], parts[2]
                        if sym not in symbols:
                            continue
                        try:
                            qty = float(parts[3])
                            px = float(parts[5]) if parts[5] else 0.0
                        except ValueError:
                            continue
                        key = (date, sym)
                        shares[key] = shares.get(key, 0.0) + qty
                        value[key] = value.get(key, 0.0) + qty * px
            return ym, shares, value
        except Exception as exc:  # noqa: BLE001
            if attempt == 3:
                print(f"  FAILED {ym}: {exc}")
            else:
                time.sleep(1.0 * (attempt + 1))
    return None


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

    prev_rows = None
    have: set[str] = set()
    if os.path.exists(OUT):
        with open(OUT, "rb") as fh:
            d = pickle.load(fh)
        prev_rows = d
        have = {str(i) for i in d["shares"].index}
    todo = [ym for ym in candidates() if ym not in have]
    if limit:
        todo = todo[:limit]
    print(f"candidates: {len(candidates())}, stored: {len(have)}, fetching: {len(todo)}")

    rows_shares: dict = {}
    rows_value: dict = {}
    with ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = {ex.submit(fetch_one, ym, symbols): ym for ym in todo}
        done = 0
        for fut in as_completed(futs):
            res = fut.result()
            done += 1
            if res is not None:
                ym, shares, value = res
                for (date, sym), q in shares.items():
                    rows_shares.setdefault(date, {})[sym] = q
                    rows_value.setdefault(date, {})[sym] = value[(date, sym)]
            if done % 24 == 0:
                print(f"  {done}/{len(todo)} files")

    def frame(rows: dict) -> pd.DataFrame:
        df = pd.DataFrame.from_dict(rows, orient="index")
        df.index = pd.to_datetime(df.index)
        return df.sort_index()

    new_shares = frame(rows_shares)
    new_value = frame(rows_value)
    parts_s = [f for f in (prev_rows["shares"] if prev_rows else None, new_shares) if f is not None and len(f)]
    parts_v = [f for f in (prev_rows["value"] if prev_rows else None, new_value) if f is not None and len(f)]
    df_s = pd.concat(parts_s).sort_index()
    df_s = df_s[~df_s.index.duplicated(keep="last")].astype("float32")
    df_v = pd.concat(parts_v).sort_index()
    df_v = df_v[~df_v.index.duplicated(keep="last")].astype("float32")
    with open(OUT, "wb") as fh:
        pickle.dump({"shares": df_s, "value": df_v}, fh)
    print(f"stored shares {df_s.shape} value {df_v.shape} -> {OUT}")


if __name__ == "__main__":
    main()