#!/usr/bin/env python3
"""Probe the top N Tranco domains for /llms.txt and AI-crawler rules in /robots.txt.

Usage: ~/.venv/bin/python probe.py [N=1000] [--summarise-only]
Raw results: ~/data/llms-txt/results.jsonl (one line per domain, resumable).
Summary: summary.md and domains.csv next to this script. Definitions: README.md.
"""
import csv
import json
import os
import re
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests

DATA = os.path.expanduser("~/data/llms-txt")
HERE = os.path.dirname(os.path.abspath(__file__))
UA = "Mozilla/5.0 (compatible; Vesper/1.0; +https://untilnextsession.com/agents)"
CRAWLERS = ["GPTBot", "ClaudeBot", "anthropic-ai", "Google-Extended", "CCBot", "PerplexityBot",
            "Bytespider", "Applebot-Extended", "meta-externalagent"]
N = int(next((a for a in sys.argv[1:] if a.isdigit()), 1000))
SUMMARISE_ONLY = "--summarise-only" in sys.argv


def get(url):
    try:
        r = requests.get(url, headers={"User-Agent": UA}, timeout=15, allow_redirects=True)
        if "charset" not in r.headers.get("content-type", "").lower():
            r.encoding = "utf-8"  # requests defaults text/* to ISO-8859-1, which garbles UTF-8 bodies
        return r.status_code, r.headers.get("content-type", ""), r.text[:200_000], None
    except Exception as e:  # noqa: BLE001
        return None, "", "", type(e).__name__


def llms_present(status, body):
    if status != 200 or not body.strip():
        return False, "no"
    head = body[:500].lower()
    if "<html" in head or "<!doctype" in head:
        return False, "html"
    first = next((ln.strip() for ln in body.splitlines() if ln.strip()), "")
    return first.startswith("#"), ("ok" if first.startswith("#") else "no-h1")


def parse_robots(body):
    """Return {agent_lower: {"disallow": [...], "allow": [...]}} merging groups per agent."""
    groups, current = {}, []
    for raw in body.splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line or ":" not in line:
            continue
        key, _, val = line.partition(":")
        key, val = key.strip().lower(), val.strip()
        if key == "user-agent":
            if current and current[-1][1]:  # a new block starts after rules were seen
                current = []
            current.append([val.lower(), False])
            groups.setdefault(val.lower(), {"disallow": [], "allow": []})
        elif key in ("disallow", "allow") and current:
            for agent in current:
                agent[1] = True
                groups[agent[0]][key].append(val)
    return groups


def blocked(groups, name):
    g = groups.get(name.lower())
    if not g:
        return "absent"
    if "/" in g["disallow"] and "/" not in g["allow"]:
        return "blocked"
    if any(d for d in g["disallow"]):
        return "restricted"
    return "allowed"


def probe(rank, domain):
    s, ct, body, err = get(f"https://{domain}/llms.txt")
    ok, why = llms_present(s, body)
    rs, rct, rbody, rerr = get(f"https://{domain}/robots.txt")
    rules = {}
    if rs == 200 and "<html" not in rbody[:500].lower():
        groups = parse_robots(rbody)
        rules = {c: blocked(groups, c) for c in CRAWLERS}
        rules["*"] = blocked(groups, "*")
    return {"rank": rank, "domain": domain, "llms_status": s, "llms_err": err, "llms_ok": ok, "llms_why": why,
            "llms_len": len(body) if ok else 0, "llms_first": (body.strip().splitlines()[0][:120] if ok else ""),
            "robots_status": rs, "robots_err": rerr, "rules": rules}


def main():
    os.makedirs(DATA, exist_ok=True)
    out = os.path.join(DATA, "results.jsonl")
    done = {}
    if os.path.exists(out):
        for ln in open(out):
            row = json.loads(ln)
            done[row["domain"]] = row
    if not SUMMARISE_ONLY:
        domains = []
        with open(os.path.join(DATA, "top-1m.csv")) as f:
            for ln in f:
                rank, dom = ln.strip().split(",", 1)
                if int(rank) > N:
                    break
                if dom not in done:
                    domains.append((int(rank), dom))
        with open(out, "a") as fo, ThreadPoolExecutor(max_workers=32) as ex:
            futs = [ex.submit(probe, r, d) for r, d in domains]
            for i, fut in enumerate(as_completed(futs), 1):
                row = fut.result()
                done[row["domain"]] = row
                fo.write(json.dumps(row) + "\n")
                fo.flush()
                if i % 100 == 0:
                    print(f"{i}/{len(domains)} probed", flush=True)
    summarise(sorted((r for r in done.values() if r["rank"] <= N), key=lambda r: r["rank"]))


def fix_mojibake(text):
    """Repair first lines captured before the charset fix: UTF-8 bytes read as Latin-1."""
    try:
        return text.encode("latin-1").decode("utf-8")
    except (UnicodeEncodeError, UnicodeDecodeError):
        return text


def summarise(rows):
    for r in rows:
        r["llms_first"] = fix_mojibake(r["llms_first"])
    with open(os.path.join(HERE, "domains.csv"), "w", newline="") as f:
        w = csv.writer(f)
        w.writerow(["rank", "domain", "llms_txt", "llms_len", "robots"] + CRAWLERS + ["*"])
        for r in rows:
            w.writerow([r["rank"], r["domain"], "yes" if r["llms_ok"] else r["llms_why"], r["llms_len"],
                        r["robots_status"] or r["robots_err"]] + [r["rules"].get(c, "") for c in CRAWLERS + ["*"]])

    def stats(sub):
        n = len(sub)
        llms = sum(r["llms_ok"] for r in sub)
        soft = sum(r["llms_why"] == "html" for r in sub)
        unreachable = sum(r["llms_status"] is None and r["robots_status"] is None for r in sub)
        robots = sum(bool(r["rules"]) for r in sub)
        any_block = sum(any(r["rules"].get(c) == "blocked" for c in CRAWLERS) for r in sub)
        any_touch = sum(any(r["rules"].get(c) in ("blocked", "restricted") for c in CRAWLERS) for r in sub)
        per = {c: sum(r["rules"].get(c) == "blocked" for r in sub) for c in CRAWLERS}
        named = {c: sum(r["rules"].get(c, "absent") != "absent" for r in sub) for c in CRAWLERS}
        reach = n - unreachable
        star = sum(r["rules"].get("*") == "blocked" for r in sub)
        either = sum(r["rules"].get("*") == "blocked" or any(r["rules"].get(c) == "blocked" for c in CRAWLERS) for r in sub)
        return n, llms, soft, unreachable, robots, any_block, any_touch, per, named, reach, star, either

    lines = ["# Results: llms.txt and AI-crawler rules on the Tranco top 1,000", "",
             "Probe run 2026-09-06 from the server; definitions in README.md; per-domain rows in domains.csv.", ""]
    for label, sub in (("Top 100", [r for r in rows if r["rank"] <= 100]), ("Top 1,000", rows)):
        n, llms, soft, unr, robots, ab, at, per, named, reach, star, either = stats(sub)
        lines += [f"## {label} (n={n}, of which {reach} answered on at least one of the two URLs)", "",
                  f"- llms.txt present and valid: **{llms}** ({llms / n:.1%} of all, {llms / reach:.1%} of those that answered); soft-404 HTML at /llms.txt: {soft}; unreachable domains: {unr}.",
                  f"- robots.txt served and parseable: {robots} ({robots / n:.1%} of all, {robots / reach:.1%} of those that answered).",
                  f"- fully block at least one named AI crawler: **{ab}** ({ab / n:.1%} of all, {ab / reach:.1%} of those that answered, {ab / max(robots, 1):.1%} of those with a robots.txt); block or restrict at least one: {at} ({at / n:.1%}).",
                  f"- block every crawler by default (`User-agent: *` with `Disallow: /`): {star} ({star / n:.1%}); block AI crawlers by name or by that default: {either} ({either / n:.1%} of all, {either / max(robots, 1):.1%} of robots.txt).", "",
                  "| crawler | named in robots.txt | fully blocked |", "|---|---|---|"]
        lines += [f"| {c} | {named[c]} ({named[c] / n:.1%}) | {per[c]} ({per[c] / n:.1%} of all, {per[c] / max(robots, 1):.1%} of robots.txt) |" for c in CRAWLERS]
        lines.append("")
    have = [r for r in rows if r["llms_ok"]]
    lines += ["## Domains with a valid llms.txt", "", "| rank | domain | bytes | first line |", "|---|---|---|---|"]
    lines += [f"| {r['rank']} | {r['domain']} | {r['llms_len']} | {r['llms_first'].replace('|', '/')} |" for r in have]
    open(os.path.join(HERE, "summary.md"), "w").write("\n".join(lines) + "\n")
    print("\n".join(lines[:22]))


if __name__ == "__main__":
    main()