#!/usr/bin/env python3
"""AI Crawler Access Index: measure how the web's most-visited sites treat AI crawlers.

For each domain: fetch /robots.txt and /llms.txt, then determine per-agent whether
the site disallows the site root for that agent.

Sample frame: Tranco top-N (research-standard, reproducible, dated permalink).
"""
import concurrent.futures as cf
import csv
import io
import json
import re
import sys
import urllib.request
import zipfile

AGENTS = [
    "GPTBot", "OAI-SearchBot", "ChatGPT-User",
    "ClaudeBot", "anthropic-ai", "Claude-Web",
    "PerplexityBot", "CCBot", "Google-Extended",
    "Bytespider", "Applebot-Extended", "meta-externalagent",
]
AGENTS_LC = [a.lower() for a in AGENTS]
UA = "Mozilla/5.0 (compatible; SuedeCrawlerIndex/1.0; +https://suedeai.ai/)"
TIMEOUT = 8


def fetch(url, max_bytes=400_000):
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
        raw = r.read(max_bytes)
        return r.status, r.headers.get("Content-Type", ""), raw


def parse_robots(text):
    """Return {agent_lc: True/False} — True means the agent is disallowed from '/'.

    Implements the grouping rule: consecutive User-agent lines share the following
    rule block. An agent's group wins over '*' when both are present.
    """
    groups = []           # list of (set_of_agents, [(directive, value), ...])
    cur_agents, cur_rules, prev_was_ua = set(), [], False
    for raw_line in text.splitlines():
        line = raw_line.split("#", 1)[0].strip()
        if not line or ":" not in line:
            continue
        field, _, value = line.partition(":")
        field = field.strip().lower()
        value = value.strip()
        if field == "user-agent":
            if not prev_was_ua and cur_agents:
                groups.append((cur_agents, cur_rules))
                cur_agents, cur_rules = set(), []
            cur_agents.add(value.lower())
            prev_was_ua = True
        elif field in ("disallow", "allow"):
            if cur_agents:
                cur_rules.append((field, value))
            prev_was_ua = False
    if cur_agents:
        groups.append((cur_agents, cur_rules))

    def blocked_by(rules):
        """Root blocked = a 'Disallow: /' with no 'Allow: /' overriding it."""
        disallow_root = any(d == "disallow" and v == "/" for d, v in rules)
        allow_root = any(d == "allow" and v == "/" for d, v in rules)
        return disallow_root and not allow_root

    star = None
    specific = {}
    for agents, rules in groups:
        for a in agents:
            if a == "*":
                star = blocked_by(rules) if star is None else (star or blocked_by(rules))
            elif a in AGENTS_LC:
                specific[a] = specific.get(a, False) or blocked_by(rules)

    out = {}
    for a in AGENTS_LC:
        if a in specific:
            out[a] = specific[a]          # named group wins
        else:
            out[a] = bool(star)           # falls back to wildcard group
    out["_named_any"] = sorted(specific.keys())
    out["_star_blocks_root"] = bool(star)
    return out


def looks_like_llms_txt(status, ctype, raw):
    if status != 200 or not raw:
        return False
    head = raw[:2000].lstrip().lower()
    if head.startswith(b"<!doctype") or head.startswith(b"<html") or b"<head" in head[:200]:
        return False
    if "html" in ctype.lower():
        return False
    return len(raw.strip()) >= 32


def probe(domain):
    rec = {"domain": domain, "robots_status": None, "llms_txt": False,
           "blocks": {}, "named_agents": [], "error": None}
    for scheme_host in (f"https://{domain}", f"https://www.{domain}"):
        try:
            status, ctype, raw = fetch(f"{scheme_host}/robots.txt")
            rec["robots_status"] = status
            if status == 200 and raw:
                parsed = parse_robots(raw.decode("utf-8", "replace"))
                rec["named_agents"] = parsed.pop("_named_any")
                parsed.pop("_star_blocks_root", None)
                rec["blocks"] = parsed
            break
        except Exception as e:
            rec["error"] = type(e).__name__
            continue
    for scheme_host in (f"https://{domain}", f"https://www.{domain}"):
        try:
            status, ctype, raw = fetch(f"{scheme_host}/llms.txt")
            if looks_like_llms_txt(status, ctype, raw):
                rec["llms_txt"] = True
            break
        except Exception:
            continue
    return rec


def load_tranco(n):
    url = "https://tranco-list.eu/top-1m.csv.zip"
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=120) as r:
        blob = r.read()
    zf = zipfile.ZipFile(io.BytesIO(blob))
    name = zf.namelist()[0]
    domains = []
    with zf.open(name) as fh:
        for row in csv.reader(io.TextIOWrapper(fh, "utf-8")):
            if len(row) >= 2:
                domains.append(row[1].strip())
            if len(domains) >= n:
                break
    return domains


def main():
    n = int(sys.argv[1]) if len(sys.argv) > 1 else 1000
    out_path = sys.argv[2] if len(sys.argv) > 2 else "crawler_index_results.json"
    domains = load_tranco(n)
    print(f"loaded {len(domains)} domains", flush=True)
    results = []
    with cf.ThreadPoolExecutor(max_workers=40) as ex:
        for i, rec in enumerate(ex.map(probe, domains), 1):
            results.append(rec)
            if i % 100 == 0:
                print(f"  {i}/{len(domains)}", flush=True)
    with open(out_path, "w") as fh:
        json.dump({"sample": "tranco-top-%d" % n, "agents": AGENTS, "results": results}, fh)
    print(f"wrote {out_path}", flush=True)


if __name__ == "__main__":
    main()