"""Draw the corpus the sealed rule describes and fetch each page once.

Population (rule.py, section THE POPULATION): feeds from research/feeds/, first
3 item links per feed in document order, <= 3 URLs per registrable domain, one
feed per domain in rank order. Bodies land in ~/data/markdown-vs-html/, which is
not committed; `index.csv` beside them records URL, domain, status, content type,
bytes, truncation and the local filename.
"""
import csv
import hashlib
import json
import os
import re
import sys
import time
import urllib.request
import urllib.error
import xml.etree.ElementTree as ET
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from urllib.parse import urljoin, urlparse

REPO = Path(__file__).resolve().parents[2]
OUT = Path.home() / "data" / "markdown-vs-html"
BODIES = OUT / "bodies"
UA = "Vesper/1.0 (+https://untilnextsession.com; research)"
CAP = 5 * 1024 * 1024
PER_DOMAIN = 3

# Registrable domain, near enough for deduplication: last two labels, with a
# short list of two-label public suffixes that appear in the top 1,000.
TWO_LABEL = {"co.uk", "com.au", "co.jp", "com.br", "co.in", "org.uk", "ac.uk",
             "gov.uk", "com.cn", "co.nz", "com.mx", "co.za", "com.tr"}


def registrable(host):
    host = (host or "").lower().split(":")[0]
    parts = [p for p in host.split(".") if p]
    if len(parts) <= 2:
        return ".".join(parts)
    if ".".join(parts[-2:]) in TWO_LABEL:
        return ".".join(parts[-3:])
    return ".".join(parts[-2:])


def get(url, cap=CAP, timeout=20):
    """One GET. Returns (status, content_type, body_bytes, final_url, truncated, error)."""
    req = urllib.request.Request(url, headers={
        "User-Agent": UA,
        "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        "Accept-Language": "en",
    })
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            body = r.read(cap + 1)
            truncated = len(body) > cap
            return (r.status, r.headers.get("Content-Type", ""), body[:cap],
                    r.url, truncated, "")
    except urllib.error.HTTPError as e:
        try:
            body = e.read(4096)
        except Exception:
            body = b""
        return (e.code, e.headers.get("Content-Type", "") if e.headers else "",
                body, url, False, "http_error")
    except Exception as e:
        return (0, "", b"", url, False, type(e).__name__ + ": " + str(e)[:120])


def feed_list():
    """Every feed the sealed rule admits, in rank order, one per domain."""
    feeds = []
    seen = set()
    with open(REPO / "research" / "feeds" / "domains.csv") as fh:
        for row in csv.DictReader(fh):
            if row["valid"] == "1" and row["first_feed_url"]:
                d = registrable(row["domain"])
                if d not in seen:
                    seen.add(d)
                    feeds.append((int(row["rank"]), d, row["first_feed_url"]))
    with open(REPO / "research" / "feeds" / "hidden.csv") as fh:
        for row in csv.DictReader(fh):
            d = registrable(row["domain"])
            if d not in seen:
                seen.add(d)
                feeds.append((10_000, d, row["feed_url"]))
    feeds.sort()
    return feeds


LINK_RE = re.compile(rb"<link[^>]*?href=[\"']([^\"']+)[\"']", re.I)


def items_from_feed(body, base):
    """First item links in DOCUMENT ORDER. RSS <item><link>, Atom <entry><link href>."""
    links = []
    try:
        root = ET.fromstring(body)
    except Exception:
        return links
    for el in root.iter():
        tag = el.tag.split("}")[-1].lower()
        if tag in ("item", "entry"):
            href = None
            for child in el:
                ctag = child.tag.split("}")[-1].lower()
                if ctag == "link":
                    rel = (child.get("rel") or "alternate").lower()
                    if child.get("href") and rel == "alternate":
                        href = child.get("href")
                        break
                    if child.text and child.text.strip():
                        href = child.text.strip()
                        break
                if ctag == "id" and href is None and (child.text or "").startswith("http"):
                    href = child.text.strip()
            if href:
                links.append(urljoin(base, href))
    return links


def main():
    BODIES.mkdir(parents=True, exist_ok=True)
    feeds = feed_list()
    print("feeds admitted by the rule: %d" % len(feeds), flush=True)

    targets = []          # (domain, feed_url, article_url)
    feed_report = []
    def pull(f):
        rank, dom, url = f
        st, ct, body, final, trunc, err = get(url, cap=2 * 1024 * 1024, timeout=20)
        links = items_from_feed(body, final) if st == 200 and body else []
        return (rank, dom, url, st, err, links)

    with ThreadPoolExecutor(max_workers=8) as ex:
        for rank, dom, url, st, err, links in ex.map(pull, feeds):
            kept = []
            for link in links:
                if not link.lower().startswith(("http://", "https://")):
                    continue
                if link in [t[2] for t in targets] or link in kept:
                    continue
                kept.append(link)
                if len(kept) >= PER_DOMAIN:
                    break
            for link in kept:
                targets.append((dom, url, link))
            feed_report.append({"rank": rank, "domain": dom, "feed": url,
                                "status": st, "error": err,
                                "items_seen": len(links), "urls_taken": len(kept)})
            print("  %-28s %s  items=%-3d taken=%d %s" %
                  (dom, st, len(links), len(kept), err), flush=True)

    print("\narticle URLs drawn: %d from %d domains" %
          (len(targets), len({t[0] for t in targets})), flush=True)

    rows = []
    def fetch_one(t):
        dom, feed, url = t
        st, ct, body, final, trunc, err = get(url)
        name = hashlib.sha256(url.encode()).hexdigest()[:16] + ".body"
        if body:
            (BODIES / name).write_bytes(body)
        return {"domain": dom, "feed": feed, "url": url, "final_url": final,
                "status": st, "content_type": ct, "bytes": len(body),
                "truncated": int(trunc), "error": err, "file": name if body else ""}

    with ThreadPoolExecutor(max_workers=6) as ex:
        for i, row in enumerate(ex.map(fetch_one, targets), 1):
            rows.append(row)
            if i % 25 == 0:
                print("  fetched %d/%d" % (i, len(targets)), flush=True)

    with open(OUT / "index.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=list(rows[0].keys()))
        w.writeheader()
        w.writerows(rows)
    (OUT / "feeds.json").write_text(json.dumps(feed_report, indent=1))

    ok = [r for r in rows if r["status"] == 200 and "html" in r["content_type"].lower()]
    print("\n%d fetched, %d are HTML 200, %d domains" %
          (len(rows), len(ok), len({r["domain"] for r in ok})))
    print("written to %s" % OUT)


if __name__ == "__main__":
    main()