"""Fetch /robots.txt for the Tranco top N with two user-agents.

Research UA (honest, names the study) and a browser-like UA. Fetching with both
lets us tell a site that genuinely has no robots.txt apart from one that simply
refuses our research agent (403/blocked), so those can be excluded rather than
counted as "no rules".
"""
import json, gzip, ssl, socket, sys, time, hashlib, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor

UA_RESEARCH = ("FentnerResearchBot/1.0 (+https://fentner.com/research; "
               "robots.txt measurement study; contact: research@fentner.com)")
UA_BROWSER = ("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 "
              "(KHTML, like Gecko) Chrome/139.0.0.0 Safari/537.36")

CTX = ssl.create_default_context(); CTX.check_hostname = False; CTX.verify_mode = ssl.CERT_NONE
socket.setdefaulttimeout(10)
MAXB = 300000


def get(url, ua):
    t0 = time.time()
    try:
        req = urllib.request.Request(url, headers={
            "User-Agent": ua, "Accept": "*/*", "Accept-Encoding": "gzip",
            "Accept-Language": "en-GB,en;q=0.9"})
        with urllib.request.urlopen(req, timeout=10, context=CTX) as x:
            raw = x.read(MAXB)
            if x.headers.get("Content-Encoding") == "gzip":
                try: raw = gzip.decompress(raw)
                except Exception: pass
            return {"status": x.status, "url": x.url,
                    "ctype": (x.headers.get("Content-Type") or "")[:120],
                    "body": raw.decode("utf-8", "replace"), "ms": int((time.time()-t0)*1000)}
    except urllib.error.HTTPError as e:
        try: b = e.read(20000).decode("utf-8", "replace")
        except Exception: b = ""
        return {"status": e.code, "url": url, "ctype": (e.headers.get("Content-Type") or "")[:120] if e.headers else "",
                "body": b, "ms": int((time.time()-t0)*1000)}
    except Exception as e:
        return {"status": -1, "url": url, "err": type(e).__name__, "ctype": "", "body": "",
                "ms": int((time.time()-t0)*1000)}


def sha(s):
    return hashlib.sha256(s.encode("utf-8", "replace")).hexdigest()[:16]


def job(t):
    rank, d = t
    u = f"https://{d}/robots.txt"
    r = get(u, UA_RESEARCH)
    if r["status"] == -1:
        u2 = f"https://www.{d}/robots.txt"
        r2 = get(u2, UA_RESEARCH)
        if r2["status"] != -1:
            u, r = u2, r2
    b = get(u, UA_BROWSER)
    out = {"rank": rank, "domain": d, "url": u,
           "r_status": r["status"], "r_url": r["url"], "r_ctype": r["ctype"], "r_err": r.get("err", ""),
           "b_status": b["status"], "b_url": b["url"], "b_ctype": b["ctype"], "b_err": b.get("err", ""),
           "r_sha": sha(r["body"]), "b_sha": sha(b["body"]),
           "r_len": len(r["body"]), "b_len": len(b["body"]),
           "r_body": r["body"]}
    out["b_body"] = "" if out["r_sha"] == out["b_sha"] else b["body"]
    return out


lo, hi = int(sys.argv[1]), int(sys.argv[2])
doms = []
for line in open(sys.argv[3] if len(sys.argv) > 3 else "../top-1m.csv"):
    rk, d = line.strip().split(",", 1); rk = int(rk)
    if lo <= rk <= hi: doms.append((rk, d))
    if rk > hi: break
print("domains:", len(doms), flush=True)
t0 = time.time()
with gzip.open(f"raw_{lo}_{hi}.jsonl.gz", "wt") as f, ThreadPoolExecutor(max_workers=20) as ex:
    for i, r in enumerate(ex.map(job, doms)):
        f.write(json.dumps(r) + "\n")
        if i % 500 == 0:
            print(i, len(doms), f"{time.time()-t0:.0f}s", flush=True)
print("done", time.time() - t0, flush=True)