"""Scan two Tranco samples for HTTP/3: Alt-Svc, DNS HTTPS record, real QUIC handshake.
Usage: python -I probe.py TRANCO_CSV OUT_JSONL [SEED]     (method: README.md; SEED default 20261028)"""
import asyncio, csv, json, random, ssl, sys, time
from concurrent.futures import ThreadPoolExecutor
from urllib.parse import urlsplit
import requests, urllib3, dns.resolver
from aioquic.asyncio import connect
from aioquic.quic.configuration import QuicConfiguration

urllib3.disable_warnings()
UA = "untilnextsession-research/1.0 (+https://untilnextsession.com; one request per site)"

def sample(path, seed=20261028):
    rows = [(int(r), d) for r, d in csv.reader(open(path))]
    top = [("top", r, d) for r, d in rows[:1000]]
    tail = random.Random(seed).sample([x for x in rows if 100_001 <= x[0] <= 1_000_000], 1000)
    return top + [("tail", r, d) for r, d in sorted(tail)]

def http(d):
    try:
        r = requests.get(f"https://{d}/", timeout=10, headers={"User-Agent": UA}, verify=False, stream=True)
        r.close()
        return {"ok": True, "status": r.status_code, "final": urlsplit(r.url).hostname,
                "alt_svc": r.headers.get("alt-svc"), "server": r.headers.get("server")}
    except Exception as e:
        return {"ok": False, "err": type(e).__name__}

def https_rr(host):
    try:
        ans = dns.resolver.resolve(host, "HTTPS", lifetime=6)
        return [a.to_text() for a in ans]
    except Exception as e:
        return type(e).__name__

async def quic(host, timeout=8.0):
    cfg = QuicConfiguration(is_client=True, alpn_protocols=["h3"], server_name=host)
    cfg.verify_mode = ssl.CERT_NONE
    try:
        async with asyncio.timeout(timeout):
            async with connect(host, 443, configuration=cfg, wait_connected=True) as c:
                return c._quic.tls.alpn_negotiated or "none"
    except BaseException as e:
        if isinstance(e, KeyboardInterrupt): raise
        return "err:" + type(e).__name__

async def quic_all(hosts):
    sem = asyncio.Semaphore(40)
    async def one(h):
        async with sem: return h, await quic(h)
    return dict(await asyncio.gather(*(one(h) for h in hosts)))

def main(src, out, seed=20261028):
    items = sample(src, seed)
    t0 = time.time()
    with ThreadPoolExecutor(48) as ex:
        hs = list(ex.map(lambda it: http(it[2]), items))
    finals = sorted({h["final"] for h in hs if h["ok"] and h["final"]})
    with ThreadPoolExecutor(32) as ex:
        rr = dict(zip(finals, ex.map(https_rr, finals)))
    q = asyncio.run(quic_all(finals))
    with open(out, "w") as f:
        for (s, r, d), h in zip(items, hs):
            fin = h.get("final")
            f.write(json.dumps({"sample": s, "rank": r, "domain": d, **h,
                                "https_rr": rr.get(fin), "quic": q.get(fin)}) + "\n")
    print(f"{len(items)} domains, {len(finals)} final hosts, {time.time()-t0:.0f} s", file=sys.stderr)

if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2], int(sys.argv[3]) if len(sys.argv) > 3 else 20261028)