"""Build the snippet index from transformers' docs/source/en/model_doc.

Run locally:      python indexer.py --out snippets.json
Push to the Hub:  python indexer.py --push your-org/transformers-snippets
(Schedule it as a daily HF Job to keep the index fresh.)
"""
import argparse
import json
import os
import re
import shutil
import subprocess
import time

REPO = "huggingface/transformers"
DOC_DIR = "docs/source/en/model_doc"

FENCE = re.compile(r"```(py|python|bash|sh)\s*\n(.*?)```", re.S)
HFOPTIONS = re.compile(r'<hfoptions id="([^"]+)">(.*?)</hfoptions>', re.S)
HFOPTION = re.compile(r'<hfoption id="([^"]+)">(.*?)</hfoption>', re.S)
HEADING = re.compile(r"^(#{2,3})\s+(.+)$", re.M)
CHECKPOINT = re.compile(
    r'(?:from_pretrained\(\s*|model\s*=\s*|--model\s+)["\']([\w.\-]+/[\w.\-]+)["\']?'
)

# v1 is inference-only: drop training and quantization snippets.
EXCLUDE = re.compile(
    r"\bTrainer\b|\.train\(|TrainingArguments|SFTTrainer|trl\b|loss\.backward|optimizer\."
    r"|BitsAndBytesConfig|TorchAoConfig|quantization_config|GPTQConfig|AwqConfig|HqqConfig"
)
INFERENCE_HINT = re.compile(r"pipeline\(|\.generate\(|from_pretrained|transformers (chat|run|serve)")


def fetch_docs(ref="main", workdir="/tmp/transformers-docs"):
    """Shallow, sparse clone of just model_doc. No GitHub API calls, so no rate limits."""
    shutil.rmtree(workdir, ignore_errors=True)
    run = lambda *a, **kw: subprocess.run(a, check=True, capture_output=True, text=True, **kw)
    run("git", "clone", "--quiet", "--depth", "1", "--filter=blob:none", "--sparse",
        "--branch", ref, f"https://github.com/{REPO}", workdir)
    run("git", "sparse-checkout", "set", DOC_DIR, cwd=workdir)
    sha = run("git", "rev-parse", "HEAD", cwd=workdir).stdout.strip()
    doc_root = os.path.join(workdir, DOC_DIR)
    files = {}
    for name in sorted(os.listdir(doc_root)):
        if name.endswith(".md"):
            with open(os.path.join(doc_root, name), encoding="utf-8") as f:
                files[f"{DOC_DIR}/{name}"] = f.read()
    return sha, files


def _section_for(text, pos):
    """Nearest heading above a position, used as context for loose code blocks."""
    title = ""
    for m in HEADING.finditer(text):
        if m.start() > pos:
            break
        title = m.group(2).strip()
    return title


def _clean(code):
    code = code.strip("\n")
    return re.sub(r"^\s*<!--.*?-->\s*$", "", code, flags=re.M).strip("\n")


def _snippet(code, lang, label, section):
    code = _clean(code)
    if not code or EXCLUDE.search(code):
        return None
    if not INFERENCE_HINT.search(code):
        return None
    lang = "python" if lang in ("py", "python") else "bash"
    ckpt = CHECKPOINT.search(code)
    return {
        "label": label,
        "section": section,
        "language": lang,
        "code": code,
        "checkpoint": ckpt.group(1) if ckpt else None,
        "needs_gpu": bool(re.search(r'cuda|device_map|attn_implementation="flash', code)),
    }


def parse_doc(text):
    snippets, covered = [], []
    # 1. Structured usage tabs: <hfoptions id="usage"><hfoption id="Pipeline">...
    for group in HFOPTIONS.finditer(text):
        covered.append(group.span())
        section = _section_for(text, group.start())
        for opt in HFOPTION.finditer(group.group(2)):
            for fence in FENCE.finditer(opt.group(2)):
                s = _snippet(fence.group(2), fence.group(1), opt.group(1), section)
                if s:
                    snippets.append(s)
    # 2. Loose code blocks under usage-like headings (older doc pages).
    for fence in FENCE.finditer(text):
        if any(a <= fence.start() < b for a, b in covered):
            continue
        section = _section_for(text, fence.start())
        if not re.search(r"usage|example|inference|generat|quick|tips|notes", section, re.I):
            continue
        s = _snippet(fence.group(2), fence.group(1), section or "Example", section)
        if s:
            snippets.append(s)
    # De-duplicate identical code.
    seen, out = set(), []
    for s in snippets:
        if s["code"] not in seen:
            seen.add(s["code"])
            out.append(s)
    return out


def build_index(ref="main"):
    sha, files = fetch_docs(ref)
    docs = {}
    for path, text in files.items():
        arch = os.path.basename(path)[:-3]
        title = re.search(r"^#\s+(.+)$", text, re.M)
        docs[arch] = {
            "title": title.group(1).strip() if title else arch,
            "doc_url": f"https://huggingface.co/docs/transformers/model_doc/{arch}",
            "source_url": f"https://github.com/{REPO}/blob/{sha}/{path}",
            "snippets": parse_doc(text),
        }
    return {"ref": ref, "sha": sha, "built_at": int(time.time()), "docs": docs}


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--ref", default="main")
    ap.add_argument("--out", default="snippets.json")
    ap.add_argument("--push", help="dataset repo id to upload snippets.json to")
    args = ap.parse_args()

    index = build_index(args.ref)
    with open(args.out, "w") as f:
        json.dump(index, f)
    n = sum(len(d["snippets"]) for d in index["docs"].values())
    print(f"{len(index['docs'])} docs, {n} snippets @ {index['sha'][:8]}")

    if args.push:
        from huggingface_hub import HfApi
        api = HfApi()
        api.create_repo(args.push, repo_type="dataset", exist_ok=True)
        api.upload_file(path_or_fileobj=args.out, path_in_repo="snippets.json",
                        repo_id=args.push, repo_type="dataset",
                        commit_message=f"Index transformers@{index['sha'][:8]}")
        print(f"Pushed to https://huggingface.co/datasets/{args.push}")