#!/usr/bin/env python3
"""Arm E source packets: union of URLs cited by arms C and D per contest -> packets/<contest>.md.

Each source's text is fetched once (HTML -> text, PDF -> pdftotext), capped per source and per packet so the
local dense model gets the same evidence the frontier models used, without search.
"""
import json, re, subprocess, sys, tempfile, urllib.request
from html.parser import HTMLParser
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parent
PER_SOURCE, PER_PACKET = 6000, 36000
URL = re.compile(r"https?://[^\s)\]>\"']+")


class Text(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.out, self.skip = [], 0

    def handle_starttag(self, tag, attrs):  # type: ignore[no-untyped-def]
        if tag in ("script", "style", "nav", "footer", "header", "noscript"):
            self.skip += 1

    def handle_endtag(self, tag):  # type: ignore[no-untyped-def]
        if tag in ("script", "style", "nav", "footer", "header", "noscript") and self.skip:
            self.skip -= 1

    def handle_data(self, d):  # type: ignore[no-untyped-def]
        if not self.skip and d.strip():
            self.out.append(d.strip())


def fetch(url: str) -> str:
    req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0 (ballot-guide research)"})
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            data, ctype = r.read(), r.headers.get("Content-Type", "")
    except Exception as e:  # noqa: BLE001
        return f"[could not fetch: {e}]"
    if "pdf" in ctype or url.lower().endswith(".pdf"):
        with tempfile.NamedTemporaryFile(suffix=".pdf") as f:
            f.write(data)
            f.flush()
            return subprocess.run(["pdftotext", "-layout", f.name, "-"], capture_output=True,
                                  text=True).stdout
    p = Text()
    p.feed(data.decode("utf-8", "replace"))
    return "\n".join(p.out)


def sources_of(page: str) -> list:
    part = re.split(r"\n#+\s*Sources", page)
    return [u.rstrip(".,;") for u in URL.findall(part[-1] if len(part) > 1 else page)]


def main() -> None:
    data = json.loads((HERE / "contests.json").read_text())
    (HERE / "packets").mkdir(exist_ok=True)
    only = set(sys.argv[1].split(",")) if len(sys.argv) > 1 else None
    for c in data["contests"]:
        if only and c["id"] not in only:
            continue
        urls = []
        for arm in ("C", "D"):
            f = ROOT / "runs" / arm / f"{c['id']}.md"
            if f.exists():
                urls += [u for u in sources_of(f.read_text()) if u not in urls]
        parts, total = [], 0
        for i, u in enumerate(urls, 1):
            t = re.sub(r"\n{3,}", "\n\n", fetch(u))[:PER_SOURCE]
            if total + len(t) > PER_PACKET:
                t = t[: max(0, PER_PACKET - total)]
            total += len(t)
            parts.append(f"### [{i}] {u}\n{t}\n")
        (HERE / "packets" / f"{c['id']}.md").write_text("\n".join(parts))
        print(c["id"], len(urls), "sources", total, "chars")


if __name__ == "__main__":
    main()