"""Fetch SEC DERA Form 13F structured data sets and aggregate per quarter.

One ZIP per quarter (2016q1..2023q4 on this host), each holding TSVs of every
13F filing parsed by the SEC (COVERPAGE + INFOTABLE). We aggregate the info
table by CUSIP: total shares, total value, number of filers, and a
representative issuer name.

    https://www.sec.gov/files/structureddata/data/form-13f-data-sets/YYYYqQ_form13f.zip

Writes data/research/dera13f/agg_YYYYqQ.csv (resume-safe: zips are cached and
quarters with an aggregate CSV are skipped) and prints a summary line per
quarter.

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

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

import pandas as pd
import requests

from research.lab.fetch_us import _proxy_env

DIR = os.path.join("data", "research", "dera13f")
BASE = "https://www.sec.gov/files/structureddata/data/form-13f-data-sets/"
UA = {"User-Agent": "quant-research/0.1 gahow@example.com"}
QUARTERS = [f"{y}q{q}" for y in range(2016, 2024) for q in (1, 2, 3, 4)]


def download(q: str) -> str | None:
    os.makedirs(DIR, exist_ok=True)
    path = os.path.join(DIR, f"{q}_form13f.zip")
    if os.path.exists(path) and os.path.getsize(path) > 1_000_000:
        return path
    url = f"{BASE}{q}_form13f.zip"
    for attempt in range(4):
        try:
            r = requests.get(url, headers=UA, timeout=600)
            r.raise_for_status()
            with open(path, "wb") as fh:
                fh.write(r.content)
            return path
        except Exception as exc:  # noqa: BLE001
            print(f"  download {q} attempt {attempt}: {exc}")
    return None


def aggregate(q: str, path: str) -> str | None:
    out = os.path.join(DIR, f"agg_{q}.csv")
    if os.path.exists(out):
        return out
    try:
        zf = zipfile.ZipFile(path)
        names = zf.namelist()
        info = next((n for n in names if "INFOTABLE" in n.upper() and n.lower().endswith(".tsv")), None)
        if info is None:
            print(f"  {q}: no INFOTABLE in {names}")
            return None
        with zf.open(info) as fh:
            df = pd.read_csv(io.TextIOWrapper(fh, encoding="latin-1"), sep="\t", dtype=str,
                             usecols=lambda c: c.upper() in
                             ("NAMEOFISSUER", "CUSIP", "VALUE", "SSHPRNAMT", "SSHPRNAMTTYPE"),
                             low_memory=False)
        df.columns = [c.upper() for c in df.columns]
        for col in ("VALUE", "SSHPRNAMT"):
            df[col] = pd.to_numeric(df[col], errors="coerce")
        df = df[df.get("SSHPRNAMTTYPE", "SH").fillna("SH").str.upper().isin(["", "SH"])]
        g = df.groupby("CUSIP", dropna=True)
        agg = pd.DataFrame({
            "issuer": g["NAMEOFISSUER"].first(),
            "shares": g["SSHPRNAMT"].sum(),
            "value": g["VALUE"].sum(),
            "n_filers": g.size(),
        }).reset_index()
        agg.to_csv(out, index=False)
        print(f"  {q}: {len(agg)} cusips, {len(df)} rows", flush=True)
        return out
    except Exception as exc:  # noqa: BLE001
        print(f"  {q}: FAILED {exc}")
        return None


def main() -> None:
    _proxy_env()
    only = sys.argv[1:] or QUARTERS
    with ThreadPoolExecutor(max_workers=3) as ex:
        futs = {ex.submit(download, q): q for q in only}
        paths = {}
        for fut in as_completed(futs):
            q = futs[fut]
            p = fut.result()
            if p:
                paths[q] = p
                print(f"  downloaded {q}", flush=True)
    for q in only:
        if q in paths:
            aggregate(q, paths[q])


if __name__ == "__main__":
    main()