"""Enrich every cited domain and every recommended vendor domain.

Cited domains  -> Tranco rank, Wayback first capture + capture count.
Vendor domains -> DNS, HTTP status, final URL after redirects.

Writes cited_domains.json and vendor_domains.json.
"""
import json, os, re, ssl, socket, sys, time, urllib.parse, urllib.request
import collections
import concurrent.futures as cf
from urllib.parse import urlparse

HERE = os.path.dirname(os.path.abspath(__file__))
UA = "Mozilla/5.0 (compatible; TrellnerResearchBot/1.0; +https://trellner.com/research)"
socket.setdefaulttimeout(10)
CTX = ssl.create_default_context()
CTX.check_hostname = False
CTX.verify_mode = ssl.CERT_NONE

RANK = {}
for line in open(os.path.join(HERE, "top-1m.csv")):
    r, d = line.strip().split(",", 1)
    RANK[d] = int(r)


def reg(host):
    """Registrable-ish domain: last two labels (good enough for .com-heavy data)."""
    parts = host.split(".")
    return ".".join(parts[-2:]) if len(parts) > 2 else host


def load():
    recs = []
    for line in open(os.path.join(HERE, "main_raw.jsonl")):
        line = line.strip()
        if line:
            recs.append(json.loads(line))
    return recs


def wayback(domain):
    """First 200-capture timestamp and total capture count for a domain."""
    out = {"first": None, "captures": 0, "wb_error": None}
    base = "http://web.archive.org/cdx/search/cdx?"
    q = {"url": domain, "matchType": "domain", "output": "json",
         "filter": "statuscode:200", "limit": "1", "fl": "timestamp"}
    for attempt in range(3):
        try:
            u = base + urllib.parse.urlencode(q)
            j = json.load(urllib.request.urlopen(
                urllib.request.Request(u, headers={"User-Agent": UA}), timeout=40))
            out["first"] = j[1][0] if len(j) > 1 else None
            break
        except Exception as e:
            out["wb_error"] = type(e).__name__
            time.sleep(3 + 3 * attempt)
    q2 = {"url": domain, "matchType": "domain", "output": "json",
          "filter": "statuscode:200", "showNumPages": "true"}
    try:
        u = base + urllib.parse.urlencode(q2)
        txt = urllib.request.urlopen(
            urllib.request.Request(u, headers={"User-Agent": UA}), timeout=40
        ).read().decode().strip()
        out["pages"] = int(txt) if txt.isdigit() else None
    except Exception:
        out["pages"] = None
    return out


def http_check(domain):
    o = {"domain": domain, "ip": None, "status": None, "final": None, "err": None}
    try:
        o["ip"] = socket.gethostbyname(domain)
    except Exception as e:
        o["dns_err"] = str(e)[:60]
        return o
    for scheme in ("https", "http"):
        try:
            req = urllib.request.Request(f"{scheme}://{domain}/", headers={"User-Agent": UA})
            r = urllib.request.urlopen(req, timeout=15, context=CTX)
            body = r.read(200000)
            o["status"] = r.status
            o["final"] = r.geturl()
            o["bytes"] = len(body)
            txt = re.sub(r"<[^>]+>", " ", re.sub(
                r"<script.*?</script>|<style.*?</style>", "",
                body.decode("utf8", "ignore"), flags=re.S))
            o["words"] = len(txt.split())
            if "DEPLOYMENT_DISABLED" in txt or "Payment required" in txt:
                o["status"] = 402
            break
        except urllib.error.HTTPError as e:
            o["status"] = e.code
            o["final"] = f"{scheme}://{domain}/"
            break
        except Exception as e:
            o["err"] = repr(e)[:90]
    return o


def main():
    recs = load()
    cited = collections.Counter()
    cited_cats = collections.defaultdict(set)
    for r in recs:
        for u in r.get("citations", []):
            h = urlparse(u).netloc.lower().replace("www.", "")
            if h:
                d = reg(h)
                cited[d] += 1
                cited_cats[d].add(r["cat"])
    print("unique cited domains:", len(cited))

    vendors = collections.Counter()
    for r in recs:
        for p in r.get("picks", []):
            d = (p.get("domain") or "").lower()
            if d and "." in d and " " not in d:
                vendors[d] += 1
    print("unique vendor domains:", len(vendors))

    # Wayback for all cited domains (politeness: 6 workers)
    doms = list(cited)
    with cf.ThreadPoolExecutor(6) as ex:
        wb = dict(zip(doms, ex.map(wayback, doms)))
    out = []
    for d in doms:
        out.append({"domain": d, "citations": cited[d],
                    "categories": len(cited_cats[d]),
                    "tranco_rank": RANK.get(d),
                    "wayback_first": wb[d]["first"],
                    "wayback_pages": wb[d].get("pages"),
                    "wayback_error": wb[d]["wb_error"]})
    json.dump(out, open(os.path.join(HERE, "cited_domains.json"), "w"), indent=1)
    print("wrote cited_domains.json")

    vd = list(vendors)
    with cf.ThreadPoolExecutor(16) as ex:
        vres = list(ex.map(http_check, vd))
    for o in vres:
        o["n"] = vendors[o["domain"]]
        o["tranco_rank"] = RANK.get(o["domain"])
    json.dump(vres, open(os.path.join(HERE, "vendor_domains.json"), "w"), indent=1)
    print("wrote vendor_domains.json")


if __name__ == "__main__":
    main()