"""Ingest SEC EDGAR filings (10-K, 10-Q, 8-K, ...) or arbitrary document URLs
(FOMC minutes, FCA handbook pages, IFRS/Basel texts) into training-ready JSONL.

Outputs plain-text chunks with metadata:
    {"text": ..., "company": ..., "form": ..., "date": ..., "source_url": ...}

These chunks are the raw material for (a) continued pretraining, or (b) synthetic
instruction generation (turning chunks into Q&A pairs with a teacher model) —
step (b) is a separate pass; this script only does clean ingestion.

SEC requires a descriptive User-Agent with contact email on all requests.

Usage:
    python -m src.data.edgar_ingest --tickers AAPL MSFT JPM --forms 10-K 10-Q --out data/edgar
    python -m src.data.edgar_ingest --url-list urls.txt --out data/regulatory
"""

import argparse
import json
import pathlib
import re
import time

import requests
from bs4 import BeautifulSoup

USER_AGENT = "FinLLM-Foundry research (finpy07@gmail.com)"
HEADERS = {"User-Agent": USER_AGENT, "Accept-Encoding": "gzip, deflate"}
RATE_LIMIT_S = 0.15  # SEC allows max 10 req/s; stay well under


def _get(url):
    time.sleep(RATE_LIMIT_S)
    r = requests.get(url, headers=HEADERS, timeout=60)
    r.raise_for_status()
    return r


def ticker_to_cik():
    data = _get("https://www.sec.gov/files/company_tickers.json").json()
    return {v["ticker"].upper(): str(v["cik_str"]).zfill(10) for v in data.values()}


def list_filings(cik, forms, limit):
    subs = _get(f"https://data.sec.gov/submissions/CIK{cik}.json").json()
    recent = subs["filings"]["recent"]
    out = []
    for form, acc, doc, date in zip(
        recent["form"], recent["accessionNumber"], recent["primaryDocument"], recent["filingDate"]
    ):
        if form in forms:
            acc_nodash = acc.replace("-", "")
            url = f"https://www.sec.gov/Archives/edgar/data/{int(cik)}/{acc_nodash}/{doc}"
            out.append({"form": form, "date": date, "url": url})
        if len(out) >= limit:
            break
    return out


def html_to_text(html):
    soup = BeautifulSoup(html, "html.parser")
    for tag in soup(["script", "style"]):
        tag.decompose()
    text = soup.get_text(separator="\n")
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[ \t]{2,}", " ", text)
    return text.strip()


def chunk(text, target_chars=4000, overlap=400):
    """Paragraph-aware chunking; ~4000 chars ≈ 1000 tokens."""
    paras = [p.strip() for p in text.split("\n\n") if len(p.strip()) > 80]
    chunks, buf = [], ""
    for p in paras:
        if len(buf) + len(p) > target_chars and buf:
            chunks.append(buf.strip())
            buf = buf[-overlap:] + "\n\n" + p
        else:
            buf += "\n\n" + p
    if len(buf.strip()) > 500:
        chunks.append(buf.strip())
    return chunks


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tickers", nargs="*", default=[])
    ap.add_argument("--forms", nargs="*", default=["10-K", "10-Q"])
    ap.add_argument("--per-company", type=int, default=4, help="filings per company")
    ap.add_argument("--url-list", default=None, help="text file, one document URL per line")
    ap.add_argument("--out", default="data/edgar")
    args = ap.parse_args()

    outdir = pathlib.Path(args.out)
    outdir.mkdir(parents=True, exist_ok=True)
    outfile = outdir / "chunks.jsonl"
    n_chunks = 0

    with open(outfile, "w") as f:
        if args.tickers:
            ciks = ticker_to_cik()
            for t in args.tickers:
                cik = ciks.get(t.upper())
                if not cik:
                    print(f"[skip] unknown ticker {t}")
                    continue
                for filing in list_filings(cik, set(args.forms), args.per_company):
                    try:
                        text = html_to_text(_get(filing["url"]).text)
                    except Exception as e:  # noqa: BLE001
                        print(f"[skip] {filing['url']}: {e}")
                        continue
                    for c in chunk(text):
                        f.write(json.dumps({"text": c, "company": t.upper(), **filing,
                                            "source_url": filing["url"]}) + "\n")
                        n_chunks += 1
                    print(f"[ok] {t} {filing['form']} {filing['date']}")

        if args.url_list:
            for url in pathlib.Path(args.url_list).read_text().splitlines():
                url = url.strip()
                if not url or url.startswith("#"):
                    continue
                try:
                    text = html_to_text(_get(url).text)
                except Exception as e:  # noqa: BLE001
                    print(f"[skip] {url}: {e}")
                    continue
                for c in chunk(text):
                    f.write(json.dumps({"text": c, "company": None, "form": "doc",
                                        "date": None, "source_url": url}) + "\n")
                    n_chunks += 1
                print(f"[ok] {url}")

    print(f"Wrote {n_chunks} chunks to {outfile}")


if __name__ == "__main__":
    main()