"""Scan the samples of rule.py. Usage: python -B scan.py TRANCO_CSV OUT_JSONL BODIES_DIR [--fixtures URL]"""
import csv, hashlib, json, os, random, sys, time
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import urlsplit
import requests, urllib3
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import rule

urllib3.disable_warnings()
UA = "untilnextsession-research/1.0 (+https://untilnextsession.com/experiments/; a handful of requests per site, once)"
THREADS = 64


def samples(path):
    rows = [(int(r), d) for r, d in csv.reader(open(path))]
    a = random.Random(rule.SEED).sample([x for x in rows if x[0] <= rule.RANK_A], rule.SIZE_A)
    b = rows[:1000]
    out = {d: {"rank": r, "a": True, "b": False} for r, d in a}
    for r, d in b:
        out.setdefault(d, {"rank": r, "a": False, "b": False})["b"] = True
    return out


def fetch(session, url):
    """(status, final url, body) or (None, None, None) when no HTTP answer came."""
    t0 = time.time()
    try:
        with session.get(url, timeout=rule.TIMEOUT_S, verify=False, stream=True,
                         headers={"User-Agent": UA, "Accept": "application/json, text/plain, */*"}) as r:
            body = b""
            for chunk in r.iter_content(16384):
                body += chunk
                if len(body) >= rule.MAX_BODY or time.time() - t0 > 2 * rule.TIMEOUT_S:
                    break
            return r.status_code, r.url, body[:rule.MAX_BODY]
    except Exception:
        return None, None, None


def scan_host(base):
    """One domain, by rule step 1. Returns None if the control got no HTTP answer."""
    s = requests.Session()
    s.max_redirects = 5
    status, final, control = fetch(s, base + rule.CONTROL)
    if status is None:
        return None
    row = {"base": base, "control": status, "paths": {}}
    for key, path in rule.PATHS.items():
        st, fin, body = fetch(s, base + path)
        ok = st is not None and rule.classify(key, st, body, control)
        row["paths"][key] = {"status": st, "ok": ok}
        if ok:
            row["paths"][key].update(final=fin, sha=hashlib.sha256(body).hexdigest()[:16], size=len(body))
            if key == "ai_catalog":
                row["paths"][key]["has_mcp"] = rule.ai_catalog_has_mcp(body)
            if key != "llms":
                row["paths"][key]["_body"] = body
    s.close()
    return row


def scan_domain(item, bodies):
    d, meta = item
    row = scan_host(f"https://{d}") or scan_host(f"https://www.{d}")
    out = {"domain": d, **meta, "reachable": row is not None}
    if row:
        for key, p in row["paths"].items():
            b = p.pop("_body", None)
            if b:
                with open(os.path.join(bodies, f"{p['sha']}.body"), "wb") as f:
                    f.write(b)
        out.update(row)
    return out


def main(src, out, bodies):
    os.makedirs(bodies, exist_ok=True)
    items = list(samples(src).items())
    t0 = time.time()
    with open(out, "w") as f, ThreadPoolExecutor(THREADS) as ex:
        for i, r in enumerate(ex.map(lambda it: scan_domain(it, bodies), items), 1):
            f.write(json.dumps(r) + "\n")
            if i % 1000 == 0:
                f.flush()
                print(f"{i}/{len(items)} {time.time()-t0:.0f} s", file=sys.stderr, flush=True)
    print(f"done {len(items)} {time.time()-t0:.0f} s", file=sys.stderr)


if __name__ == "__main__":
    main(*sys.argv[1:4])