#!/usr/bin/env python3
"""aicheck - can AI search engines read your site?

Checks one website and reports whether the crawlers behind ChatGPT, Perplexity
and Claude are allowed to read it, whether the homepage ships readable text
without JavaScript, and whether an llms.txt exists.

Single file, standard library only, no signup, no network calls except to the
site you name.

    python aicheck.py example.com
    python aicheck.py example.com --json
    python aicheck.py example.com --all

MIT licensed. Built by Reese Calder.
https://ai-visibility.lastminutedealshq.com
"""

import argparse
import gzip
import json
import re
import sys
import time
import urllib.error
import urllib.request

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except Exception:
    pass

__version__ = "1.0.0"

UA = "aicheck/%s (+https://ai-visibility.lastminutedealshq.com/tools/aicheck.py)" % __version__

# agent -> (display name, vendor, job, why blocking it matters)
AGENTS = [
    ("oai-searchbot", "OAI-SearchBot", "OpenAI", "search",
     "Indexes pages for ChatGPT Search. Blocking it costs you ChatGPT citations."),
    ("perplexitybot", "PerplexityBot", "Perplexity", "search",
     "Indexes pages for Perplexity answers and citations."),
    ("claude-searchbot", "Claude-SearchBot", "Anthropic", "search",
     "Indexes pages so Claude can cite them."),
    ("gptbot", "GPTBot", "OpenAI", "training",
     "OpenAI's model-training crawler. Blocking it does not affect ChatGPT Search."),
    ("claudebot", "ClaudeBot", "Anthropic", "training",
     "Anthropic's model-training crawler."),
    ("google-extended", "Google-Extended", "Google", "training",
     "Only controls Gemini training. Does not affect Google Search or AI Overviews."),
    ("ccbot", "CCBot", "Common Crawl", "training",
     "Common Crawl, an open dataset many models train on."),
    ("chatgpt-user", "ChatGPT-User", "OpenAI", "user-action",
     "Fetches your page live when a user asks ChatGPT about it."),
    ("perplexity-user", "Perplexity-User", "Perplexity", "user-action",
     "Fetches your page live when a user asks Perplexity about it."),
    ("googlebot", "Googlebot", "Google", "classic-search",
     "Ordinary Google Search, including AI Overviews. Shown for comparison."),
]

JOB_LABEL = {
    "search": "AI search and citations",
    "training": "Model training",
    "user-action": "Live user fetches",
    "classic-search": "Classic search (for comparison)",
}


_LINES = re.compile(r"\r\n|\r|\n")  # robots.cc line terminators: LF, CR, CRLF

def fetch(url, timeout=10, limit=600000):
    """Return (status, text). status is an int, or a string for local failures."""
    req = urllib.request.Request(url, headers={
        "User-Agent": UA,
        "Accept": "text/html,text/plain,*/*",
        "Accept-Encoding": "gzip",
    })
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            raw = r.read(limit)
            if r.headers.get("Content-Encoding") == "gzip":
                try:
                    raw = gzip.decompress(raw)
                except Exception:
                    pass
            return (r.status, raw.decode("utf-8", "replace"))
    except urllib.error.HTTPError as e:
        return (e.code, "")
    except Exception as e:
        return (type(e).__name__, "")


def looks_html(body):
    head = body.lstrip()[:200].lower()
    return head.startswith("<!doctype") or head.startswith("<html") or head.startswith("<head")


def fetch_robots_retrying(origin, timeout=10):
    """Some sites' CDN caches a bad redirect-to-homepage response for the exact
    /robots.txt path (verified live against nytimes.com's Fastly edge, which serves
    a stale redirect for the plain path while any URL with a query string is a cache
    miss that reaches origin and gets the real file). Reporting unknown off one
    cache-poisoned fetch is a false negative for a real, readable file, so retry once
    with a cache-busting query string before giving up."""
    status, txt = fetch(origin + "/robots.txt", timeout)
    bad = (status == 200 and looks_html(txt)) or (isinstance(status, int) and status >= 500)
    if not bad:
        return status, txt
    return fetch(origin + "/robots.txt?_cb=" + str(int(time.time() * 1000)), timeout)


def parse_groups(txt):
    """Parse robots.txt into [{agents, rules}]. Consecutive User-agent lines
    share one rule block, which is what the standard requires."""
    groups, cur, expecting = [], None, False
    # robots.cc treats BOTH 0x0A and 0x0D as line terminators (CRLF counts as
    # one), so a file with bare CR endings is many lines to Google and ONE line
    # to a split("\n"), which swallows every Disallow and fails open (alrosa.ru).
    for raw in _LINES.split(txt):
        line = raw.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 cur is not None and expecting:
                cur["agents"].append(value.lower())
            else:
                cur = {"agents": [value.lower()], "rules": []}
                groups.append(cur)
            expecting = True
        elif field in ("allow", "disallow"):
            if cur is None:
                cur = {"agents": ["*"], "rules": []}
                groups.append(cur)
            expecting = False
            cur["rules"].append((field == "allow", value))
    return groups


def _matches(path, pattern):
    """Port of RobotsMatchStrategy::Matches from Google's reference robots.txt
    parser (github.com/google/robotstxt). The pattern is anchored at the start
    of the path, "*" is a wildcard, and "$" is special only as the final
    character. Without this, "Disallow: *" reads as "matches nothing" and the
    check fails open, and "Allow: /$" is missed entirely."""
    pathlen = len(path)
    pos = [0] * (pathlen + 1)
    numpos = 1
    n = len(pattern)
    for i in range(n):
        ch = pattern[i]
        if ch == "$" and i + 1 == n:
            return pos[numpos - 1] == pathlen
        if ch == "*":
            numpos = pathlen - pos[0] + 1
            for j in range(1, numpos):
                pos[j] = pos[j - 1] + 1
        else:
            newnumpos = 0
            for j in range(numpos):
                if pos[j] < pathlen and path[pos[j]] == ch:
                    pos[newnumpos] = pos[j] + 1
                    newnumpos += 1
            numpos = newnumpos
            if numpos == 0:
                return False
    return True



def product_token(value):
    """The matchable part of a robots.txt user-agent value.

    Port of RobotsMatcher::ExtractUserAgent from Google's reference parser: the
    token runs while characters are [a-zA-Z_-] and stops at the first character
    outside that set. RFC 9309's product-token grammar is the same, and it has no
    digits, so "AI2Bot" is read as "AI", "MJ12bot" as "MJ" and "Kangaroo Bot" as
    "Kangaroo". The comparison that follows is EXACT and case-insensitive, not a
    prefix test: a "User-agent: Google" group does not apply to Google-Extended.
    """
    out = []
    for ch in value:
        if ch.isalpha() or ch in "-_":
            out.append(ch)
        else:
            break
    return "".join(out)


def blocked_for(groups, agent, path="/"):
    """Is `agent` disallowed from `path`?

    Follows Google's reference matcher: EVERY group naming the crawler
    contributes rules (they are merged, not just the first match), the longest
    matching pattern wins, Allow wins ties, and a pattern's priority is its
    length, so a bare "Disallow:" scores 0 and is ignored. If any
    crawler-specific group matched, the catch-all groups are not consulted.

    Returns (blocked, named) where `named` says robots.txt mentioned this
    crawler by name rather than falling back to the catch-all group.
    """
    specific, star = [], []
    for g in groups:
        ags = g["agents"]
        # RFC 9309: a group applies when its product token matches our crawler's
        # token, treating the robots.txt token as a prefix. Matching in the other
        # direction as well (agent in a) is wrong: it makes a "User-agent:
        # Googlebot-News" group apply to Googlebot and report it as blocked when
        # it is not. News sites commonly carry such groups.
        if any(a != "*" and product_token(a).lower() == agent.lower() for a in ags):
            specific.append(g)
        elif "*" in ags:
            star.append(g)
    grps = specific or star
    if not grps:
        return (False, False)
    allow = disallow = 0
    for g in grps:
        for is_allow, pat in g["rules"]:
            if not _matches(path, pat):
                continue
            prio = len(pat)
            if is_allow:
                allow = max(allow, prio)
            else:
                disallow = max(disallow, prio)
    if allow > 0 or disallow > 0:
        return (disallow > allow, bool(specific))
    return (False, bool(specific))

def readable_check(status, body):
    """How much text a crawler that does not run JavaScript actually gets."""
    # Same rule as robots.txt: a homepage we could not fetch is unknown, not unreadable.
    # Some sites reject non-browser clients outright, and calling that "not readable"
    # would be a false negative.
    if not isinstance(status, int) or status >= 400 or not body.strip():
        return {"readable": None, "text_chars": 0, "script_tags": 0,
                "note": "Could not fetch the homepage (returned %s), so I cannot tell how much "
                        "text a crawler would get." % status}
    scripts = len(re.findall(r"<script\b", body, re.I))
    txt = re.sub(r"(?is)<(script|style|noscript|template)\b.*?</\1>", " ", body)
    txt = re.sub(r"(?s)<[^>]+>", " ", txt)
    txt = re.sub(r"\s+", " ", txt).strip()
    chars = len(txt)
    readable = chars >= 600
    if readable:
        note = "The page ships real text in its HTML, so an AI crawler can read it."
    elif scripts >= 3:
        note = ("The HTML has almost no text and several scripts. This looks like a JavaScript "
                "app that renders in the browser, which most AI crawlers cannot see. "
                "Server-side or static rendering fixes it.")
    else:
        note = ("The page returned very little text. Worth checking that real content is in the "
                "HTML, not loaded in later by JavaScript.")
    return {"readable": readable, "text_chars": chars, "script_tags": scripts, "note": note}


def check(site, timeout=10):
    site = site.strip()
    if not re.match(r"^https?://", site):
        site = "https://" + site
    m = re.match(r"^(https?://[^/]+)", site)
    if not m:
        return {"error": "That does not look like a valid website address."}
    origin = m.group(1)

    r_status, r_txt = fetch_robots_retrying(origin, timeout)
    # Some sites serve an HTML page instead of a real robots.txt
    if isinstance(r_status, int) and r_status == 200 and looks_html(r_txt):
        r_status, r_txt = ("html-not-robots", "")

    # Three states. A missing robots.txt (404, or an empty 200) really does mean
    # crawling is allowed. A fetch that failed tells us nothing, and answering
    # "allowed" for it would be a false all-clear, so that case reports as unknown.
    robots_ok = isinstance(r_status, int) and 200 <= r_status < 300
    no_robots = r_status == 404 or (robots_ok and not r_txt.strip())
    robots_unknown = not robots_ok and r_status != 404
    groups = parse_groups(r_txt) if (robots_ok and r_txt.strip()) else []

    bots = []
    for agent, name, vendor, job, why in AGENTS:
        blocked, named = blocked_for(groups, agent)
        if robots_unknown:
            allowed_state = None
        elif no_robots:
            allowed_state = True
        else:
            allowed_state = not blocked
        bots.append({"ua": name, "vendor": vendor, "job": job, "why": why,
                     "allowed": allowed_state, "named_explicitly": named})

    h_status, h_body = fetch(origin + "/", timeout)
    render = readable_check(h_status, h_body)

    l_status, l_body = fetch(origin + "/llms.txt", timeout)
    has_llms = (isinstance(l_status, int) and l_status == 200
                and bool(l_body.strip()) and not looks_html(l_body))

    blocked_search = [b["ua"] for b in bots if b["job"] == "search" and b["allowed"] is False]
    problems = []
    if robots_unknown:
        problems.append("Could not read %s/robots.txt (it returned %s), so the crawler results "
                        "are unknown rather than allowed." % (origin, r_status))
    if blocked_search:
        problems.append("Blocked from AI answers: your robots.txt blocks " +
                        ", ".join(blocked_search) + ".")
    if render["readable"] is False:
        problems.append("Your homepage may not be readable by AI crawlers "
                        "(little or no text in the HTML).")
    if render["readable"] is None:
        problems.append("Could not fetch the homepage, so I cannot tell whether it sends "
                        "readable text. Some sites reject non-browser clients.")
    blocked_training = [b["ua"] for b in bots if b["job"] == "training" and b["allowed"] is False]
    if blocked_training and not blocked_search:
        problems.append("You block training crawlers (" + ", ".join(blocked_training) +
                        ") but still allow the search crawlers, so citations are unaffected.")

    if robots_unknown:
        verdict = ("Could not read this site's robots.txt, so I cannot say whether AI search "
                   "crawlers are allowed.")
    elif blocked_search or render["readable"] is False:
        verdict = "There are things holding your AI visibility back."
    elif render["readable"] is None:
        verdict = "Could not fetch this site's homepage, so this check is incomplete."
    else:
        verdict = ("Looking good. AI search crawlers can reach you and your homepage is readable."
                   + (" You also have an llms.txt." if has_llms else ""))

    return {
        "site": origin,
        "verdict": verdict,
        "problems": problems,
        "checks": {
            "crawler_access": {
                "robots_status": str(r_status),
                "robots_readable": not robots_unknown,
                "bots": bots,
                "blocked_search_count": len(blocked_search),
            },
            "readable_content": render,
            "llms_txt": {"present": has_llms, "status": str(l_status)},
        },
    }


def render_text(d, show_all=False):
    if "error" in d:
        return d["error"]
    out = []
    out.append("")
    out.append(d["site"])
    out.append("=" * min(len(d["site"]), 60))
    out.append("")
    bots = d["checks"]["crawler_access"]["bots"]
    jobs = ["search", "training", "user-action", "classic-search"] if show_all else ["search"]
    for job in jobs:
        rows = [b for b in bots if b["job"] == job]
        if not rows:
            continue
        out.append("  " + JOB_LABEL[job])
        for b in rows:
            if b["allowed"] is None:
                state = "unknown"
                named = ""
            else:
                state = "allowed" if b["allowed"] else "BLOCKED"
                named = "" if b["named_explicitly"] else "  (via catch-all rule)"
            out.append("    %-9s %-18s %s%s" % (state, b["ua"], b["vendor"], named))
        out.append("")
    rc = d["checks"]["readable_content"]
    if rc["readable"] is None:
        out.append("  Homepage readable without JavaScript: unknown")
    else:
        out.append("  Homepage readable without JavaScript: %s (%s characters of text)"
                   % ("yes" if rc["readable"] else "NO", format(rc["text_chars"], ",")))
    if rc["readable"] is not True:
        out.append("    " + rc["note"])
    out.append("  llms.txt: %s" % ("found" if d["checks"]["llms_txt"]["present"] else "not found"))
    out.append("")
    out.append("  " + d["verdict"])
    for p in d["problems"]:
        out.append("  - " + p)
    out.append("")
    if not show_all:
        out.append("  Run with --all to see training and user-action crawlers too.")
        out.append("")
    return "\n".join(out)


def main():
    ap = argparse.ArgumentParser(
        prog="aicheck",
        description="Check whether AI search engines can read a website.")
    ap.add_argument("site", help="domain or URL, for example example.com")
    ap.add_argument("--json", action="store_true", dest="as_json",
                    help="print the full result as JSON")
    ap.add_argument("--all", action="store_true", dest="show_all",
                    help="show training and user-action crawlers as well")
    ap.add_argument("--timeout", type=float, default=10, help="per-request timeout in seconds")
    ap.add_argument("--version", action="version", version="aicheck " + __version__)
    a = ap.parse_args()

    d = check(a.site, timeout=a.timeout)
    if a.as_json:
        print(json.dumps(d, indent=1))
    else:
        print(render_text(d, show_all=a.show_all))
    if "error" in d:
        return 2
    # exit codes, so this is usable in CI:
    #   0 = AI search crawlers can reach the site
    #   1 = at least one AI search crawler is blocked
    #   2 = the address could not be parsed
    #   3 = the robots.txt could not be read, so the answer is unknown
    if not d["checks"]["crawler_access"]["robots_readable"]:
        return 3
    return 1 if d["checks"]["crawler_access"]["blocked_search_count"] else 0


if __name__ == "__main__":
    sys.exit(main())