"""Fetch every Form 990-PF the IRS has published for one private foundation, straight from irs.gov.
Standard library only. Nothing is asked of ParetoAlpha and nothing is sent to it.

    python3 irs990pf_fetch.py 38-2966992 --out returns/        # every e-filed 990-PF for the EIN, as <OBJECT_ID>.xml
    python3 verify_chain.py --read returns/                     # then read them in order, as a carryover chain

How it finds them, all from U.S. Government public files:
  1. index_YYYY.csv (apps.irs.gov/pub/epostcard/990/xml/YYYY/): every e-filed exempt-organization return the IRS
     published that year, with its EIN, form, tax period and OBJECT_ID. Downloaded once per year into --cache.
  2. The year's zips, as listed on irs.gov/charities-non-profits/form-990-series-downloads. A return is located by its
     member name (<OBJECT_ID>_public.xml) in each zip's central directory, which is read with two Range requests and
     kept in --cache; the zip itself (up to several GB) is never downloaded. A return is usually in a zip of its index
     year and sometimes in the year beside it, so both are searched.
  3. The one member, read with a single Range request and inflated here (Deflate64 included, which Python's zipfile
     cannot do), its CRC-32 checked against the zip's own record.

Two things the downloads page does not tell you (found 2026-09-29): the IRS server holds a second 2022 zip the page
does not list (2022_TEOS_XML_02A.zip), and the 2020 index points at returns that sit in 2019's zips.

The first run for a year downloads that year's index (50 to 95 MB) and central directories; later runs reuse them.
"""
import argparse
import csv
import gzip
import os
import re
import struct
import sys
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from verify_proof import member_from_local, LOCAL_SLACK  # the Range read's inflater and CRC check, beside this file

BASE = "https://apps.irs.gov/pub/epostcard/990/xml"
DOWNLOADS = "https://www.irs.gov/charities-non-profits/form-990-series-downloads"
UNLISTED = {"2022": [f"{BASE}/2022/2022_TEOS_XML_02A.zip"]}
UA = {"User-Agent": "irs990pf_fetch.py (research; standard library)"}


def get(url, start=None, end=None, timeout=300):
    headers = dict(UA)
    if start is not None:
        headers["Range"] = f"bytes={start}-{end - 1}"
    with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=timeout) as r:
        return r.read()


def size_of(url):
    req = urllib.request.Request(url, method="HEAD", headers=UA)
    with urllib.request.urlopen(req, timeout=60) as r:
        return int(r.headers["Content-Length"])


def zip_list(cache):
    path = cache / "zips.txt"
    if path.exists():
        urls = path.read_text().split()
    else:
        html = get(DOWNLOADS, timeout=60).decode("utf-8", "replace")
        urls = sorted(set(re.findall(r"https://apps\.irs\.gov/pub/epostcard/990/xml/\d{4}/[^\"\s<]+\.zip", html)))
        path.write_text("\n".join(urls) + "\n")
    by = {}
    for u in urls:
        by.setdefault(u.split("/")[-2], []).append(u)
    for y, extra in UNLISTED.items():
        for u in extra:
            if u not in by.setdefault(y, []):
                by[y].append(u)
    return by


def central_directory(url, cache):
    """name -> (header offset, compressed size, size, crc, method), from the zip's central directory (ZIP64 honoured)."""
    path = cache / (url.split("/")[-1] + ".cd.gz")
    if path.exists():
        cd = gzip.decompress(path.read_bytes())
    else:
        total = size_of(url)
        tail_len = min(total, 65557 + 20 + 56)
        tail_start = total - tail_len
        tail = get(url, tail_start, total)
        i = tail.rfind(b"PK\x05\x06")
        if i < 0:
            raise ValueError(f"{url}: no end-of-central-directory record")
        entries, size, offset = struct.unpack_from("<H", tail, i + 10)[0], struct.unpack_from("<I", tail, i + 12)[0], struct.unpack_from("<I", tail, i + 16)[0]
        if entries == 0xFFFF or size == 0xFFFFFFFF or offset == 0xFFFFFFFF:
            loc = i - 20
            if loc < 0 or tail[loc:loc + 4] != b"PK\x06\x07":
                raise ValueError(f"{url}: ZIP64 fields without a locator")
            rec = struct.unpack_from("<Q", tail, loc + 8)[0] - tail_start
            size, offset = struct.unpack_from("<Q", tail, rec + 40)[0], struct.unpack_from("<Q", tail, rec + 48)[0]
        cd = get(url, offset, offset + size)
        path.write_bytes(gzip.compress(cd, 6))
    out, p = {}, 0
    while p + 46 <= len(cd) and cd[p:p + 4] == b"PK\x01\x02":
        method, = struct.unpack_from("<H", cd, p + 10)
        crc, csize, usize, nlen, elen, clen = struct.unpack_from("<IIIHHH", cd, p + 16)
        off, = struct.unpack_from("<I", cd, p + 42)
        name = cd[p + 46:p + 46 + nlen].decode("utf-8", "replace")
        e, stop = p + 46 + nlen, p + 46 + nlen + elen
        while e + 4 <= stop:
            eid, elen2 = struct.unpack_from("<HH", cd, e)
            if eid == 1:
                q = e + 4
                if usize == 0xFFFFFFFF:
                    usize, = struct.unpack_from("<Q", cd, q); q += 8
                if csize == 0xFFFFFFFF:
                    csize, = struct.unpack_from("<Q", cd, q); q += 8
                if off == 0xFFFFFFFF:
                    off, = struct.unpack_from("<Q", cd, q); q += 8
            e += 4 + elen2
        out[name.split("/")[-1]] = (off, csize, usize, crc, method)
        p = stop + clen
    return out


def index_rows(year, ein, cache):
    path = cache / f"index_{year}.csv"
    if not path.exists():
        print(f"  downloading index_{year}.csv …", file=sys.stderr)
        try:
            path.write_bytes(get(f"{BASE}/{year}/index_{year}.csv"))
        except Exception as e:  # a year the IRS has not published yet
            print(f"  index_{year}.csv: {e}", file=sys.stderr)
            return []
    with open(path, newline="", encoding="utf-8", errors="replace") as f:
        return [(r["OBJECT_ID"], r["TAX_PERIOD"]) for r in csv.DictReader(f) if r.get("EIN") == ein and r.get("RETURN_TYPE") == "990PF"]


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("ein", help="the foundation's EIN, with or without the dash")
    ap.add_argument("--out", default="returns", help="folder for <OBJECT_ID>.xml (default: returns/)")
    ap.add_argument("--cache", default=str(Path.home() / ".cache" / "irs990pf"), help="where indexes and central directories are kept")
    ap.add_argument("--years", default="2019-2026", help="index years to read, e.g. 2021-2026")
    a = ap.parse_args()
    ein = re.sub(r"[\s-]", "", a.ein)
    if not re.fullmatch(r"\d{9}", ein):
        sys.exit("an EIN is nine digits")
    lo, hi = (int(x) for x in a.years.split("-"))
    cache, out = Path(a.cache), Path(a.out)
    cache.mkdir(parents=True, exist_ok=True)
    out.mkdir(parents=True, exist_ok=True)
    zips = zip_list(cache)

    wanted = {}
    for y in range(lo, hi + 1):
        for object_id, period in index_rows(y, ein, cache):
            wanted[object_id] = (y, period)
    if not wanted:
        sys.exit(f"the IRS indexes for {lo} to {hi} list no e-filed Form 990-PF for EIN {ein}")
    print(f"EIN {ein}: {len(wanted)} e-filed Form 990-PF return(s) in the indexes for {lo} to {hi}")

    got = 0
    for object_id, (y, period) in sorted(wanted.items(), key=lambda kv: kv[1][1]):
        target = out / f"{object_id}.xml"
        if target.exists():
            print(f"  {period}  {object_id}  already in {out}/")
            got += 1
            continue
        name = f"{object_id}_public.xml"
        found = None
        for year in (y, y - 1, y + 1):
            for url in zips.get(str(year), []):
                try:
                    entry = central_directory(url, cache).get(name)
                except Exception as e:
                    print(f"  {url.split('/')[-1]}: {e}", file=sys.stderr)
                    continue
                if entry:
                    found = (url, entry)
                    break
            if found:
                break
        if not found:
            print(f"  {period}  {object_id}  in no zip the IRS has published")
            continue
        url, (off, csize, usize, crc, method) = found
        raw = member_from_local(get(url, off, off + LOCAL_SLACK + csize), {"method": method, "crc": crc, "compressedSize": csize, "size": usize})
        target.write_bytes(raw)
        got += 1
        print(f"  {period}  {object_id}  {len(raw):>8,} bytes  from {url.split('/')[-1]}")
    print(f"{got} of {len(wanted)} written to {out}/. Next: python3 verify_chain.py --read {out}")
    return 0 if got else 1


if __name__ == "__main__":
    sys.exit(main())