"""Collect the raw data for the AGENTS.md study (Stride research, 2026-09).

Population: every public GitHub repository with at least MIN_STARS stars
that is not a fork, not archived, and received a push in the PUSH_WINDOW_DAYS
days before the snapshot. It is a census of that population, not a sample.

Pipeline (each stage checkpoints to OUT, so a rerun resumes):
  1. enumerate   GitHub repository search, split into star ranges small
                 enough that no range exceeds the search API's 1,000-result
                 cap, paged by stars. Deduplicated by repository id.
  2. files       GraphQL, batched: the root tree (names + modes, to find
                 case variants and symlinks) plus every instruction file
                 this study tracks, with text for AGENTS.md and CLAUDE.md.
  3. variants    Text for AGENTS.md / CLAUDE.md files whose root name is a
                 case variant (Agents.md, claude.md, ...).
  4. history     For repositories with an AGENTS.md or CLAUDE.md: commits
                 touching the file on the default branch (count, and the
                 dates of up to 100, to date first appearance).
  5. extras      Symlinked AGENTS.md / CLAUDE.md files followed to the file
                 they point at, and .claude/CLAUDE.md (entries, text,
                 history) for every repository with a .claude directory.

Raw output (repo metadata and file TEXT) stays local and is never
published: the files belong to their repositories. analyze.py derives the
per-repository metrics that are published.

Usage: python collect.py --out DIR [--snapshot YYYY-MM-DD] [--stages files,variants,history,extras]
Auth: GITHUB_TOKEN, or the GitHub CLI's token (`gh auth token`); read-only
public data. Python 3.10+, standard library only.
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path

MIN_STARS = 5000
PUSH_WINDOW_DAYS = 90
SEARCH_CAP = 1000
# Split star ranges well below the cap: counts drift while we page (a repo
# gains stars or a push), and a range that grows past 1,000 mid-crawl would
# silently lose its tail.
SPLIT_ABOVE = 900
BATCH = 20

# Instruction files checked in every repository. Root files are also matched
# case-insensitively against the root tree, so "Agents.md" still counts.
ROOT_FILES = {
    "agents_md": "AGENTS.md",
    "claude_md": "CLAUDE.md",
    "gemini_md": "GEMINI.md",
    "agent_md": "AGENT.md",
    "cursorrules": ".cursorrules",
    "windsurfrules": ".windsurfrules",
    "clinerules": ".clinerules",
}
NESTED = {
    "copilot_instructions": ".github/copilot-instructions.md",
    "cursor_rules_dir": ".cursor/rules",
    "windsurf_rules_dir": ".windsurf/rules",
    "junie_guidelines": ".junie/guidelines.md",
}
TEXT_KEYS = ("agents_md", "claude_md")


def token() -> str:
    """A GitHub token: GITHUB_TOKEN if set, else the GitHub CLI's (`gh auth
    token`; set GH_BIN when gh is not on PATH). Read-only public data needs
    no scopes."""
    if os.environ.get("GITHUB_TOKEN"):
        return os.environ["GITHUB_TOKEN"]
    gh = os.environ.get("GH_BIN", "gh")
    return subprocess.run([gh, "auth", "token"], capture_output=True, text=True, check=True).stdout.strip()


TOKEN = token()


def http(method: str, url: str, body: dict | None = None, attempt: int = 0) -> dict:
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method, headers={
        "Authorization": f"Bearer {TOKEN}",
        "Accept": "application/vnd.github+json",
        "User-Agent": "stride-research-agents-md",
        "Content-Type": "application/json",
    })
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        payload = e.read().decode(errors="replace")
        # Secondary rate limits and transient 5xx: back off and retry.
        if e.code in (403, 429, 500, 502, 503, 504) and attempt < 8:
            wait = int(e.headers.get("Retry-After") or 0) or min(60 * (attempt + 1), 300)
            print(f"  HTTP {e.code}, retry in {wait}s: {payload[:160]}", flush=True)
            time.sleep(wait)
            return http(method, url, body, attempt + 1)
        raise RuntimeError(f"HTTP {e.code} for {url}: {payload[:500]}") from e
    except (urllib.error.URLError, TimeoutError) as e:
        if attempt < 8:
            time.sleep(15 * (attempt + 1))
            return http(method, url, body, attempt + 1)
        raise


def search(q: str, page: int = 1, per_page: int = 100) -> dict:
    url = "https://api.github.com/search/repositories?" + urllib.parse.urlencode(
        {"q": q, "sort": "stars", "order": "desc", "per_page": per_page, "page": page})
    time.sleep(2.2)  # search API: 30 requests/minute
    return http("GET", url)


def graphql(query: str) -> dict:
    out = http("POST", "https://api.github.com/graphql", {"query": query})
    if out.get("errors") and not out.get("data"):
        raise RuntimeError(json.dumps(out["errors"])[:800])
    return out


# ---------------------------------------------------------------- stage 1

def enumerate_population(out: Path, snapshot: dt.date) -> list[dict]:
    path = out / "population.json"
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))["repos"]
    since = (snapshot - dt.timedelta(days=PUSH_WINDOW_DAYS)).isoformat()
    base = f"fork:false archived:false pushed:>={since}"
    total = search(f"stars:>={MIN_STARS} {base}", per_page=1)["total_count"]
    print(f"population total_count={total}", flush=True)

    # Split [lo, hi] star ranges until each holds <= SEARCH_CAP repositories.
    todo, ranges = [(MIN_STARS, 10_000_000)], []
    while todo:
        lo, hi = todo.pop()
        n = search(f"stars:{lo}..{hi} {base}", per_page=1)["total_count"]
        if n <= SPLIT_ABOVE or lo == hi:
            ranges.append((lo, hi, n))
        else:
            mid = (lo + hi) // 2 if hi < 10_000_000 else max(lo * 2, lo + 1)
            mid = min(mid, hi - 1)
            todo += [(lo, mid), (mid + 1, hi)]
    ranges.sort()
    print(f"{len(ranges)} star ranges", flush=True)

    repos: dict[int, dict] = {}
    for lo, hi, n in ranges:
        if n > SEARCH_CAP:
            print(f"  WARNING range {lo}..{hi} holds {n} > cap; truncated", flush=True)
        pages = min((n + 99) // 100, SEARCH_CAP // 100)
        for page in range(1, pages + 1):
            res = search(f"stars:{lo}..{hi} {base}", page=page)
            if res.get("incomplete_results"):
                print(f"  incomplete results {lo}..{hi} p{page}; retrying once", flush=True)
                res = search(f"stars:{lo}..{hi} {base}", page=page)
            for item in res["items"]:
                repos[item["id"]] = {
                    "id": item["id"],
                    "full_name": item["full_name"],
                    "stars": item["stargazers_count"],
                    "language": item["language"],
                    "pushed_at": item["pushed_at"],
                    "created_at": item["created_at"],
                    "default_branch": item["default_branch"],
                    "fork": item["fork"],
                    "archived": item["archived"],
                    "topics": item.get("topics", []),
                }
        print(f"  {lo}..{hi}: {n} expected, {len(repos)} collected so far", flush=True)

    data = {
        "snapshot": snapshot.isoformat(),
        "query": f"stars:>={MIN_STARS} {base}",
        "search_total_count": total,
        "ranges": ranges,
        "repos": sorted(repos.values(), key=lambda r: -r["stars"]),
    }
    path.write_text(json.dumps(data, indent=1), encoding="utf-8")
    return data["repos"]


# ---------------------------------------------------------------- stage 2

BLOB_TEXT = "... on Blob { byteSize isBinary oid text }"
BLOB_META = "... on Blob { byteSize isBinary oid }"


def repo_fragment(alias: str, full_name: str) -> str:
    owner, name = full_name.split("/", 1)
    parts = [f'{alias}: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{',
             "databaseId nameWithOwner stargazerCount isFork isArchived pushedAt",
             "primaryLanguage { name } defaultBranchRef { name }",
             'root: object(expression: "HEAD:") { ... on Tree { entries { name type mode } } }']
    for key, p in ROOT_FILES.items():
        frag = BLOB_TEXT if key in TEXT_KEYS else BLOB_META
        parts.append(f'{key}: object(expression: {json.dumps("HEAD:" + p)}) {{ {frag} }}')
    for key, p in NESTED.items():
        if p.endswith("rules"):
            parts.append(f'{key}: object(expression: {json.dumps("HEAD:" + p)}) {{ ... on Tree {{ entries {{ name type }} }} ... on Blob {{ byteSize }} }}')
        else:
            parts.append(f'{key}: object(expression: {json.dumps("HEAD:" + p)}) {{ {BLOB_META} }}')
    parts.append("}")
    return "\n".join(parts)


def fetch_files(out: Path, repos: list[dict]) -> dict:
    path = out / "files.json"
    done: dict = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    todo = [r for r in repos if r["full_name"] not in done]
    print(f"files: {len(done)} done, {len(todo)} to fetch", flush=True)
    for i in range(0, len(todo), BATCH):
        chunk = todo[i:i + BATCH]
        body = "\n".join(repo_fragment(f"r{j}", r["full_name"]) for j, r in enumerate(chunk))
        query = "query {\n rateLimit { cost remaining resetAt }\n" + body + "\n}"
        try:
            res = graphql(query)
        except RuntimeError as e:
            # One bad repository (e.g. a timeout on a huge root tree) should
            # not sink the batch: retry its members one by one.
            print(f"  batch failed ({e}); retrying singly", flush=True)
            for r in chunk:
                try:
                    one = graphql("query {\n" + repo_fragment("r0", r["full_name"]) + "\n}")
                    done[r["full_name"]] = one.get("data", {}).get("r0")
                except RuntimeError as e2:
                    done[r["full_name"]] = {"error": str(e2)[:300]}
            continue
        data = res.get("data") or {}
        for j, r in enumerate(chunk):
            done[r["full_name"]] = data.get(f"r{j}")  # None when renamed/removed since search
        rl = data.get("rateLimit") or {}
        if (i // BATCH) % 10 == 0:
            path.write_text(json.dumps(done), encoding="utf-8")
            print(f"  {len(done)}/{len(repos)} cost={rl.get('cost')} remaining={rl.get('remaining')}", flush=True)
        if rl.get("remaining", 5000) < 200:
            reset = dt.datetime.fromisoformat(rl["resetAt"].replace("Z", "+00:00"))
            wait = max(0, (reset - dt.datetime.now(dt.timezone.utc)).total_seconds()) + 5
            print(f"  graphql budget low; sleeping {wait:.0f}s", flush=True)
            time.sleep(wait)
    path.write_text(json.dumps(done), encoding="utf-8")
    return done


# ---------------------------------------------------------------- stage 3

def root_name(entries: list[dict] | None, canonical: str) -> dict | None:
    for e in entries or []:
        if e["name"].lower() == canonical.lower():
            return e
    return None


def fetch_variants(out: Path, files: dict) -> dict:
    """Text for AGENTS.md / CLAUDE.md present at the root under another case."""
    path = out / "variants.json"
    done: dict = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    wanted = []
    for full, rec in files.items():
        if not rec or "error" in rec:
            continue
        entries = (rec.get("root") or {}).get("entries")
        for key in TEXT_KEYS:
            e = root_name(entries, ROOT_FILES[key])
            if e and e["name"] != ROOT_FILES[key] and f"{full}::{key}" not in done:
                wanted.append((full, key, e["name"]))
    print(f"variants: {len(wanted)} case-variant files to fetch", flush=True)
    for full, key, actual in wanted:
        owner, name = full.split("/", 1)
        q = (f'query {{ r: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{ '
             f'f: object(expression: {json.dumps("HEAD:" + actual)}) {{ {BLOB_TEXT} }} }} }}')
        res = graphql(q)
        done[f"{full}::{key}"] = {"name": actual, "blob": ((res.get("data") or {}).get("r") or {}).get("f")}
    path.write_text(json.dumps(done), encoding="utf-8")
    return done


# ---------------------------------------------------------------- stage 4

def fetch_history(out: Path, files: dict, variants: dict) -> dict:
    path = out / "history.json"
    done: dict = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
    wanted = []
    for full, rec in files.items():
        if not rec or "error" in rec:
            continue
        entries = (rec.get("root") or {}).get("entries")
        for key in TEXT_KEYS:
            e = root_name(entries, ROOT_FILES[key])
            if e and f"{full}::{key}" not in done:
                wanted.append((full, key, e["name"]))
    print(f"history: {len(wanted)} files", flush=True)
    step = 10
    for i in range(0, len(wanted), step):
        chunk = wanted[i:i + step]
        frags = []
        for j, (full, key, actual) in enumerate(chunk):
            owner, name = full.split("/", 1)
            frags.append(
                f'h{j}: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{ '
                f'defaultBranchRef {{ target {{ ... on Commit {{ history(first: 100, path: {json.dumps(actual)}) '
                f'{{ totalCount nodes {{ committedDate }} }} }} }} }} }}')
        res = graphql("query {\n rateLimit { remaining resetAt }\n" + "\n".join(frags) + "\n}")
        data = res.get("data") or {}
        for j, (full, key, actual) in enumerate(chunk):
            node = data.get(f"h{j}") or {}
            hist = (((node.get("defaultBranchRef") or {}).get("target") or {}).get("history")) or {}
            done[f"{full}::{key}"] = {
                "total": hist.get("totalCount"),
                "dates": [n["committedDate"] for n in hist.get("nodes") or []],
            }
        if (i // step) % 20 == 0:
            path.write_text(json.dumps(done), encoding="utf-8")
            print(f"  {len(done)}/{len(wanted)} remaining={(data.get('rateLimit') or {}).get('remaining')}", flush=True)
    path.write_text(json.dumps(done), encoding="utf-8")
    return done


# ---------------------------------------------------------------- stage 5

SYMLINK_MODE = 0o120000  # git's mode for a symbolic link


def resolve_path(base_dir: str, target: str) -> str | None:
    """A symlink's target, resolved against the link's directory, as a path
    from the repository root. None when it would leave the repository."""
    parts = [p for p in base_dir.split("/") if p]
    for seg in target.strip().replace("\\", "/").split("/"):
        if seg in ("", "."):
            continue
        if seg == "..":
            if not parts:
                return None
            parts.pop()
        else:
            parts.append(seg)
    return "/".join(parts) or None


def repo_query(full_name: str, body: str) -> dict:
    owner, name = full_name.split("/", 1)
    q = f"query {{ r: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{ {body} }} }}"
    return (graphql(q).get("data") or {}).get("r") or {}


def follow_link(full_name: str, path: str, max_hops: int = 5) -> dict:
    """Follow symlinks from `path` to a regular file and return its text.

    A symlink's blob holds its target path, so a CLAUDE.md that links to
    AGENTS.md reads as the string "AGENTS.md". Content metrics need the
    file the link points at."""
    hops: list[str] = []
    cur = path
    for _ in range(max_hops):
        parent, _, base = cur.rpartition("/")
        r = repo_query(full_name,
                       f'd: object(expression: {json.dumps("HEAD:" + parent)}) {{ ... on Tree {{ entries {{ name mode }} }} }} '
                       f'f: object(expression: {json.dumps("HEAD:" + cur)}) {{ {BLOB_TEXT} }}')
        entry = next((e for e in ((r.get("d") or {}).get("entries") or []) if e["name"] == base), None)
        blob = r.get("f")
        if entry is None or blob is None:
            return {"path": cur, "hops": hops, "text": None, "missing": True}
        if entry["mode"] != SYMLINK_MODE:
            return {"path": cur, "hops": hops, "text": blob.get("text"), "byteSize": blob.get("byteSize")}
        hops.append(cur)
        nxt = resolve_path(parent, blob.get("text") or "")
        if not nxt:
            return {"path": None, "hops": hops, "text": None, "escapes": True}
        cur = nxt
    return {"path": cur, "hops": hops, "text": None, "tooManyHops": True}


def fetch_extras(out: Path, files: dict) -> dict:
    """Two things stages 2-4 do not cover.

    1. Symlink targets. When AGENTS.md or CLAUDE.md is a symlink, follow it
       to the file it points at (the other file, or a third one such as
       .ai/AGENTS.md or CONTRIBUTING.md).
    2. .claude/CLAUDE.md. Claude Code reads a project's CLAUDE.md from
       either ./CLAUDE.md or ./.claude/CLAUDE.md, and either one stops it
       reading AGENTS.md by default (code.claude.com/docs/en/memory). For
       every repository with a .claude directory: its entries, the text of
       .claude/CLAUDE.md (following a symlink), and that file's history.
    """
    path = out / "extras.json"
    done: dict = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}

    def save() -> None:
        path.write_text(json.dumps(done), encoding="utf-8")

    readable = {full: rec for full, rec in files.items() if rec and "error" not in rec}

    # 1. Root symlinks.
    links = []
    for full, rec in readable.items():
        entries = (rec.get("root") or {}).get("entries")
        for key in TEXT_KEYS:
            e = root_name(entries, ROOT_FILES[key])
            if e and e.get("mode") == SYMLINK_MODE and f"{full}::{key}::link" not in done:
                links.append((full, key, e["name"]))
    print(f"extras: {len(links)} root symlinks to follow", flush=True)
    for full, key, actual in links:
        done[f"{full}::{key}::link"] = follow_link(full, actual)
    save()

    # 2a. .claude directories, batched.
    dirs = [full for full, rec in readable.items()
            if any(e["name"] == ".claude" and e["type"] == "tree"
                   for e in ((rec.get("root") or {}).get("entries") or []))
            and f"{full}::dotclaude" not in done]
    print(f"extras: {len(dirs)} .claude directories", flush=True)
    for i in range(0, len(dirs), BATCH):
        chunk = dirs[i:i + BATCH]
        frags = []
        for j, full in enumerate(chunk):
            owner, name = full.split("/", 1)
            frags.append(
                f'r{j}: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{ '
                f't: object(expression: "HEAD:.claude") {{ ... on Tree {{ entries {{ name type mode }} }} }} '
                f'c: object(expression: "HEAD:.claude/CLAUDE.md") {{ {BLOB_TEXT} }} }}')
        data = graphql("query {\n" + "\n".join(frags) + "\n}").get("data") or {}
        for j, full in enumerate(chunk):
            node = data.get(f"r{j}") or {}
            entries = (node.get("t") or {}).get("entries") or []
            done[f"{full}::dotclaude"] = {"entries": entries, "claude_md": node.get("c")}
        save()

    # 2b. .claude/CLAUDE.md under another case, or as a symlink: follow it.
    for full in [f for f in readable if f"{f}::dotclaude" in done]:
        rec = done[f"{full}::dotclaude"]
        e = next((x for x in rec["entries"] if x["name"].lower() == "claude.md"), None)
        if not e or f"{full}::dotclaude::link" in done:
            continue
        if e["name"] != "CLAUDE.md" or e.get("mode") == SYMLINK_MODE:
            done[f"{full}::dotclaude::link"] = follow_link(full, ".claude/" + e["name"])
    save()

    # 2c. History of .claude/CLAUDE.md, as stage 4 does for the root files.
    wanted = []
    for full in readable:
        rec = done.get(f"{full}::dotclaude")
        e = next((x for x in (rec or {}).get("entries", []) if x["name"].lower() == "claude.md"), None)
        if e and f"{full}::dotclaude::history" not in done:
            wanted.append((full, ".claude/" + e["name"]))
    print(f"extras: history for {len(wanted)} .claude/CLAUDE.md files", flush=True)
    for i in range(0, len(wanted), 10):
        chunk = wanted[i:i + 10]
        frags = []
        for j, (full, p) in enumerate(chunk):
            owner, name = full.split("/", 1)
            frags.append(
                f'h{j}: repository(owner: {json.dumps(owner)}, name: {json.dumps(name)}) {{ '
                f'defaultBranchRef {{ target {{ ... on Commit {{ history(first: 100, path: {json.dumps(p)}) '
                f'{{ totalCount nodes {{ committedDate }} }} }} }} }} }}')
        data = graphql("query {\n" + "\n".join(frags) + "\n}").get("data") or {}
        for j, (full, p) in enumerate(chunk):
            node = data.get(f"h{j}") or {}
            hist = (((node.get("defaultBranchRef") or {}).get("target") or {}).get("history")) or {}
            done[f"{full}::dotclaude::history"] = {
                "total": hist.get("totalCount"),
                "dates": [n["committedDate"] for n in hist.get("nodes") or []],
            }
        save()
    return done


STAGES = ("files", "variants", "history", "extras")


def main() -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--snapshot", default=dt.date.today().isoformat())
    ap.add_argument("--stages", default=",".join(STAGES),
                    help="comma-separated subset of " + ",".join(STAGES) + " (the population is always loaded)")
    args = ap.parse_args()
    stages = set(args.stages.split(","))
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    snapshot = dt.date.fromisoformat(args.snapshot)
    repos = enumerate_population(out, snapshot)
    print(f"population: {len(repos)} repositories", flush=True)
    files_path = out / "files.json"
    files = fetch_files(out, repos) if "files" in stages else json.loads(files_path.read_text(encoding="utf-8"))
    variants = fetch_variants(out, files) if "variants" in stages else {}
    if "history" in stages:
        fetch_history(out, files, variants)
    if "extras" in stages:
        fetch_extras(out, files)
    print("done", flush=True)


if __name__ == "__main__":
    main()