# /// script
# requires-python = ">=3.10"
# dependencies = ["pandas", "pyarrow", "huggingface_hub>=0.30", "langdetect"]
# ///
"""Turn XiaomiMiMo/MiMo-V2.6-RL-oss into the static data the explorer loads.

    uv run python -m app.mimo.build_data            # downloads what it needs from the Hub
    RAW_DIR=raw uv run python -m app.mimo.build_data  # or reuse a local copy of the same files

Writes app/mimo/data/, gzipped (Spaces serve static files uncompressed, so the page inflates them itself):
    index.json.gz      one small record per environment: what search, filters and the map need
    <domain>.json.gz   full briefs and domain-specific detail, fetched when an environment is opened
    rewards.json.gz    counts behind the Reward design page (how General's rubrics are built)

    uv run python -m app.mimo.build_data --rewards  # only rewards.json.gz (needs just the rubric files)

Rubric *questions* are kept; `pass_anchor` and `gold_answer` are not. They are the answers,
and a browsing tool is not the place to hand them out.
"""

from __future__ import annotations

import collections
import gzip
import json
import os
import re
import sys
from pathlib import Path

import pandas as pd

REPO = "XiaomiMiMo/MiMo-V2.6-RL-oss"
OUT = Path(__file__).parent / "data"
PARQUETS = {
    "code": "code.parquet",
    "webdev": "webdev.parquet",
    "cyber": "cyber.parquet",
    "music": "music.parquet",
    "general": "general/train.parquet",
}


# ── fetch ────────────────────────────────────────────────────────────────────
def fetch() -> tuple[Path, list[str]]:
    from huggingface_hub import HfApi, snapshot_download

    files = HfApi().list_repo_files(REPO, repo_type="dataset")
    raw = os.environ.get("RAW_DIR")
    if raw:
        return Path(raw), files
    path = snapshot_download(
        REPO,
        repo_type="dataset",
        allow_patterns=list(PARQUETS.values()) + ["general/envs/*/verifier_meta.json"],
    )
    return Path(path), files


# ── small helpers ────────────────────────────────────────────────────────────
def clip(text: str, n: int) -> str:
    text = re.sub(r"\s+", " ", text).strip()
    if len(text) <= n:
        return text
    cut = text[: n - 1]
    if " " in cut[n // 2:]:            # end on a whole word when there is one to end on
        cut = cut[: cut.rfind(" ")]
    return cut.rstrip(" ,;:-") + "…"


def plain(text: str) -> str:
    """Markdown down to prose, for one-line snippets: no headings, emphasis, fences or links."""
    text = re.sub(r"```.*?```", " ", text, flags=re.S)
    text = re.sub(r"^\s*#{1,6}\s*", "", text, flags=re.M)
    text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", text)
    text = re.sub(r"(\*\*|__|\*|`)", "", text)
    return re.sub(r"^\s*[-*•]\s+", "", text, flags=re.M)


def snippet(text: str, title: str, n: int = 280) -> str:
    """The brief after its headline, so a card does not say the same sentence twice."""
    body = re.sub(r"\s+", " ", plain(text)).strip()
    stem = title.rstrip("…").strip()
    at = body.find(stem) if stem else -1
    if 0 <= at < 60:                       # the headline, possibly after a greeting
        end = at + len(stem)
        if title.endswith("…"):            # it was cut: carry on from the end of that sentence
            m = re.search(r"[.!?。！？](\s|$)", body[end:])
            end = end + m.end() if m else len(body)
        rest = body[end:].lstrip(" .,;:!?。，；：！？")
        if len(rest) > 30:
            body = rest
    return clip(body, n)


def first_sentence(text: str, n: int = 110) -> str:
    text = re.sub(r"^\s*(prompt|task)\s*:\s*", "", text.strip(), flags=re.I)
    lines = [l.strip().lstrip("#").strip().strip("*").strip() for l in text.splitlines()]
    lines = [l for l in lines if l]
    # skip greetings and stubs ("Hi,", "Hello team") in favour of the first line that says something
    line = next((l for l in lines if len(l) >= 18 and not re.match(r"(?i)^(hi|hello|hey|dear|greetings)\b", l)),
                lines[0] if lines else text)
    # end of the first sentence, not the dot in "U.S." / "Inc." / "e.g."
    end = None
    for m in re.finditer(r"[.!?。！？](?=\s|$)", line):
        before = line[:m.start()].rsplit(None, 1)[-1] if line[:m.start()].strip() else ""
        if m.group() == "." and (re.fullmatch(r"(?:[A-Za-z]\.)+[A-Za-z]?", before) or before.lower() in ABBREV or len(before) == 1):
            continue
        if m.end() > 25:
            end = m.end()
            break
    return clip(line[:end] if end else line, n)


ABBREV = {"inc", "corp", "co", "ltd", "llc", "plc", "no", "mr", "ms", "mrs", "dr", "st", "vs", "etc", "jr", "sr", "fig", "approx"}


def lang_of(text: str) -> str:
    """A coarse language label. CJK is decided by script, the rest by langdetect."""
    if re.search(r"[一-鿿]", text):
        return "zh"
    from langdetect import DetectorFactory, detect

    DetectorFactory.seed = 0
    try:
        return detect(text[:600])
    except Exception:
        return "unknown"


LANG_NAMES = {
    "en": "English", "zh": "Chinese", "es": "Spanish", "tr": "Turkish", "ar": "Arabic",
    "fr": "French", "de": "German", "pt": "Portuguese", "ru": "Russian", "it": "Italian",
    "ja": "Japanese", "ko": "Korean", "id": "Indonesian", "vi": "Vietnamese", "nl": "Dutch",
    "pl": "Polish", "hi": "Hindi", "fa": "Persian", "uk": "Ukrainian", "th": "Thai",
}


def lang_name(code: str) -> str:
    return LANG_NAMES.get(code, "Other")


def instance(row) -> dict:
    return json.loads(row["extra_info"]["instance_json"])


def brief(row) -> str:
    return row["prompt"][-1]["content"]


# ── code ─────────────────────────────────────────────────────────────────────
EXT_LANG = {
    "py": "Python", "go": "Go", "js": "JavaScript", "jsx": "JavaScript", "mjs": "JavaScript",
    "cjs": "JavaScript", "ts": "TypeScript", "tsx": "TypeScript", "rb": "Ruby", "php": "PHP",
    "java": "Java", "rs": "Rust", "c": "C", "h": "C", "cpp": "C++", "cc": "C++", "hpp": "C++",
    "cs": "C#", "kt": "Kotlin", "scala": "Scala", "swift": "Swift", "dart": "Dart", "lua": "Lua",
    "ex": "Elixir", "exs": "Elixir", "hs": "Haskell", "vue": "Vue", "svelte": "Svelte",
}


def build_code(df):
    index, detail = [], {}
    for _, row in df.iterrows():
        j = instance(row)
        text = brief(row)
        # patches mix `+++ b/path` and bare `+++ path` headers; /dev/null is a deletion, not a file
        touched = [f for f in re.findall(r"^\+\+\+ (?:b/)?(\S+)", j.get("test_patch") or "", re.M) if f != "/dev/null"]
        langs = collections.Counter(
            EXT_LANG[f.rsplit(".", 1)[-1].lower()] for f in touched if f.rsplit(".", 1)[-1].lower() in EXT_LANG
        )
        language = langs.most_common(1)[0][0] if langs else "Unknown"
        eid = j["instance_id"]
        title = first_sentence(re.sub(r"^\[[A-Z]+\]\s*", "", text))
        index.append({
            "id": eid, "d": "code", "t": title,
            "s": snippet(re.sub(r"^\[[A-Z]+\]\s*", "", text), title), "f": {"prog": language},
        })
        detail[eid] = {
            "brief": text,
            "meta": [["Programming language", language], ["Test command", j.get("test_command", "")],
                     ["Working directory", j.get("cwd", "")], ["Verifier timeout", f"{j.get('verifier_timeout_sec', '?')}s"]],
            "files": {"Files the tests touch": touched[:40]},
        }
    return index, detail


# ── cyber ────────────────────────────────────────────────────────────────────
CRASH = re.compile(r"^(?P<san>[\w ]*?Sanitizer|[\w-]+):\s*(?P<kind>[\w-]+)(?: on [\w-]+)?\s+in\s+(?:function\s+)?`?(?P<fn>[^`\s]+)`?(?:.*?in file `(?P<file>[^`]+)`)?")


def build_cyber(df):
    index, detail = [], {}
    for _, row in df.iterrows():
        j = instance(row)
        text = brief(row)
        m = CRASH.match(text.strip())
        san = m.group("san") if m else "Unknown"
        kind = m.group("kind") if m else "unknown"
        fn = m.group("fn") if m else ""
        file = (m.group("file") if m else "") or ""
        project = file.split("/")[0] if "/" in file else "unknown"
        eid = j["instance_id"]
        index.append({
            "id": eid, "d": "cyber", "t": f"{kind} in {fn}" if fn else clip(text, 110),
            "s": clip(plain(text), 280),
            "f": {"crash": kind, "project": project, "sanitizer": san},
        })
        detail[eid] = {
            "brief": text,
            "meta": [["Crash", kind], ["Sanitizer", san], ["Project", project], ["Function", fn],
                     ["File", file], ["Source", "ARVO (reproduced OSS-Fuzz crashes)"]],
            "extra": clip(j.get("description", ""), 2000) if j.get("description") != text else "",
        }
    return index, detail


# ── webdev ───────────────────────────────────────────────────────────────────
# Specific kinds of site. "Landing page" is a format more than a kind, so it only wins when
# nothing more specific is named; otherwise the earliest specific mention is the primary type.
SITE_TYPES = [
    ("Restaurant & food", r"restaurant|\bcaf[eé]\b|bakery|bistro|coffee shop|food truck|\bdiner\b|pizzeria"),
    ("Portfolio", r"portfolio"), ("Dashboard", r"dashboard"),
    ("SaaS / product", r"\bsaas\b|product (?:page|showcase|site|launch)"),
    ("Blog & editorial", r"\bblog\b|magazine|newsletter|news site"),
    ("E-commerce", r"e-?commerce|online store|storefront|\bshop\b"),
    ("Agency", r"\bagency\b"), ("Game", r"\bgame\b"),
    ("Event", r"event (?:landing|site|page|website|microsite)|conference|festival|wedding"),
    ("Tool / app", r"\bweb app\b|calculator|\btracker\b|\bquiz\b"),
]
LANDING = r"landing page"
FRAMEWORKS = [
    ("React", r"\breact\b"), ("Vue", r"\bvue\b"), ("Next.js", r"next\.?js"), ("Astro", r"\bastro\b"),
    ("Svelte", r"svelte"), ("Three.js", r"three\.?js|webgl"), ("Tailwind", r"tailwind"),
    ("Plain HTML/CSS", r"plain html|vanilla (?:js|javascript)|html,? css"),
]
STYLES = [
    ("Minimalist", r"minimalis|\bminimal\b"), ("Dark mode", r"dark[- ]mode|dark theme"), ("Brutalist", r"brutalis"),
    ("Glassmorphism", r"glassmorph|glass cards|frosted glass"), ("Neumorphism", r"neumorph"), ("Swiss", r"\bswiss\b"),
    ("Scandinavian", r"scandinav|nordic"), ("Art deco", r"art[- ]deco"), ("Retro", r"\bretro\b|vintage|\by2k\b"),
    ("Cyberpunk", r"cyberpunk|\bneon\b"), ("Material", r"material design"), ("Flat", r"flat design"),
    ("Luxury", r"luxur"), ("Playful", r"playful|whimsical|cartoon"), ("Corporate", r"corporate"),
]


def primary_site(text: str) -> str:
    low = text.lower()
    hits = [(m.start(), name) for name, pat in SITE_TYPES if (m := re.search(pat, low))]
    if hits:
        return min(hits)[1]
    return "Landing page" if re.search(LANDING, low) else "Other"


def tags(text: str, vocab) -> list[str]:
    low = text.lower()
    return [name for name, pat in vocab if re.search(pat, low)]


def build_webdev(df):
    index, detail = [], {}
    for _, row in df.iterrows():
        j = instance(row)
        text = brief(row)
        site = primary_site(text)
        fws = tags(text, FRAMEWORKS) or ["Unspecified"]
        styles = tags(text, STYLES)
        language = lang_name(lang_of(text))
        eid = j["instance_id"]
        index.append({
            "id": eid, "d": "webdev", "t": first_sentence(text), "s": snippet(text, first_sentence(text)),
            "f": {"site": site, "framework": fws, "style": styles or ["None named"], "language": language},
        })
        detail[eid] = {
            "brief": text,
            "meta": [["Site type", site], ["Framework (mentioned)", ", ".join(fws)],
                     ["Style (mentioned)", ", ".join(styles) or "none named"], ["Brief language", language]],
        }
    return index, detail


# ── music ────────────────────────────────────────────────────────────────────
# The dataset's style tags are Chinese on every row, including English briefs.
MUSIC_STYLE = {
    "马林巴曲": ("Marimba piece", "Solo instrument"), "爱尔兰 reel": ("Irish reel", "Folk & world"),
    "谐谑曲": ("Scherzo", "Classical forms"), "八音盒音乐": ("Music box", "Solo instrument"),
    "小夜曲": ("Serenade", "Classical forms"), "马祖卡": ("Mazurka", "Dance"), "钢琴独奏": ("Piano solo", "Piano"),
    "弦乐合奏": ("String ensemble", "Ensemble"), "波罗乃兹": ("Polonaise", "Dance"),
    "管风琴前奏曲": ("Organ prelude", "Classical forms"), "木管五重奏": ("Woodwind quintet", "Ensemble"),
    "极简主义钢琴循环（Glass 风）": ("Minimalist piano loop (Glass)", "Piano"),
    "阿尔贝蒂低音钢琴小品": ("Alberti-bass piano piece", "Piano"), "简单钢琴小品": ("Simple piano piece", "Piano"),
    "长笛独奏": ("Flute solo", "Solo instrument"), "Stride 钢琴": ("Stride piano", "Jazz & blues"),
    "古典小步舞曲": ("Classical minuet", "Dance"), "巴赫风格四声部众赞歌": ("Bach-style chorale", "Classical forms"),
    "拉格泰姆钢琴": ("Ragtime piano", "Jazz & blues"), "莫扎特风格钢琴小品": ("Mozart-style piano piece", "Piano"),
    "手风琴曲": ("Accordion piece", "Solo instrument"), "肖邦风格圆舞曲": ("Chopin-style waltz", "Dance"),
    "回旋曲": ("Rondo", "Classical forms"), "铜管五重奏": ("Brass quintet", "Ensemble"),
    "弦乐四重奏（海顿风）": ("String quartet (Haydn)", "Ensemble"), "萨蒂风格裸体歌舞": ("Satie-style gymnopédie", "Piano"),
    "波尔卡": ("Polka", "Dance"), "英式乡村舞曲": ("English country dance", "Dance"),
    "木管三重奏": ("Woodwind trio", "Ensemble"), "新古典钢琴（Einaudi 风）": ("Neoclassical piano (Einaudi)", "Piano"),
    "卡农": ("Canon", "Classical forms"), "俄罗斯民歌": ("Russian folk song", "Folk & world"),
    "巴洛克二声部创意曲": ("Baroque two-part invention", "Classical forms"), "爱尔兰 jig": ("Irish jig", "Folk & world"),
    "竖琴曲": ("Harp piece", "Solo instrument"), "古典主题与变奏": ("Theme and variations", "Classical forms"),
    "古典小奏鸣曲": ("Classical sonatina", "Classical forms"), "桑巴": ("Samba", "Folk & world"),
    "摇摆爵士大乐队": ("Swing big band", "Jazz & blues"), "Afrobeats 现代": ("Afrobeats", "Pop & modern"),
    "巴尔干 7/8 舞曲": ("Balkan 7/8 dance", "Folk & world"), "灵歌": ("Spiritual", "Folk & world"),
    "苏格兰 hornpipe": ("Scottish hornpipe", "Folk & world"), "古筝独奏": ("Guzheng solo", "Chinese traditional"),
    "12 小节布鲁斯钢琴": ("12-bar blues piano", "Jazz & blues"), "管弦乐团全奏": ("Full orchestra", "Ensemble"),
    "Reggaeton dembow": ("Reggaeton dembow", "Pop & modern"), "口琴曲": ("Harmonica piece", "Solo instrument"),
    "比波普": ("Bebop", "Jazz & blues"), "浪漫派大提琴旋律": ("Romantic cello melody", "Solo instrument"),
    "J-Pop": ("J-Pop", "Pop & modern"), "浪漫派小提琴小品": ("Romantic violin piece", "Solo instrument"),
    "拉丁爵士": ("Latin jazz", "Jazz & blues"), "钢琴四手联弹": ("Piano four hands", "Piano"),
    "氛围钢琴": ("Ambient piano", "Piano"), "西部乡村": ("Country & western", "Folk & world"),
    "琵琶曲": ("Pipa piece", "Chinese traditional"), "萨克斯独奏": ("Saxophone solo", "Solo instrument"),
    "中国风流行": ("Chinese-style pop", "Pop & modern"), "江南丝竹": ("Jiangnan sizhu", "Chinese traditional"),
    "尤克里里曲": ("Ukulele piece", "Solo instrument"), "小号独奏": ("Trumpet solo", "Solo instrument"),
    "木吉他弹唱伴奏": ("Acoustic guitar accompaniment", "Pop & modern"), "摇滚乐队编制": ("Rock band", "Pop & modern"),
    "新灵魂乐": ("Neo-soul", "Pop & modern"), "维也纳圆舞曲": ("Viennese waltz", "Dance"),
    "肖邦风格夜曲": ("Chopin-style nocturne", "Piano"), "蓝草乐": ("Bluegrass", "Folk & world"),
    "爵士钢琴三重奏": ("Jazz piano trio", "Jazz & blues"), "克莱兹默旋律": ("Klezmer melody", "Folk & world"),
    "R&B 慢歌": ("R&B ballad", "Pop & modern"), "小提琴独奏": ("Violin solo", "Solo instrument"),
    "合成器浪潮（Synthwave）": ("Synthwave", "Pop & modern"), "德彪西风格印象派钢琴": ("Debussy-style impressionist piano", "Piano"),
    "吉他指弹": ("Fingerstyle guitar", "Solo instrument"), "爵士华尔兹": ("Jazz waltz", "Jazz & blues"),
    "大提琴独奏": ("Cello solo", "Solo instrument"), "City Pop（日式）": ("City pop", "Pop & modern"),
    "竹笛曲": ("Dizi (bamboo flute) piece", "Chinese traditional"),
}


def music_key(text: str) -> str:
    m = re.search(r"([升降]?)([A-G])([#♯b♭]?)(大调|小调)", text)
    if m:
        acc = {"升": "♯", "#": "♯", "♯": "♯", "降": "♭", "b": "♭", "♭": "♭"}.get(m.group(1) or m.group(3), "")
        return f"{m.group(2)}{acc} {'major' if m.group(4) == '大调' else 'minor'}"
    m = re.search(r"\b([A-G])(?:([#♯b♭])| (sharp|flat))? (major|minor)\b", text)
    if m:
        acc = "♯" if (m.group(2) in ("#", "♯") or m.group(3) == "sharp") else "♭" if (m.group(2) in ("b", "♭") or m.group(3) == "flat") else ""
        return f"{m.group(1)}{acc} {m.group(4)}"
    return ""


def tempo(bpm) -> str:
    try:
        b = float(bpm)
    except (TypeError, ValueError):
        return "Unknown"
    return "Slow (<80)" if b < 80 else "Moderate (80–120)" if b <= 120 else "Fast (>120)"


def build_music(df):
    index, detail = [], {}
    for n, (_, row) in enumerate(df.iterrows()):
        e = row["extra_info"]
        text = brief(row)
        eid = f"music-{e.get('src_id', n)}"
        voices = int(e.get("nvoice_want") or 0)
        tag = str(e.get("tag") or "")
        style, family = MUSIC_STYLE.get(tag, (tag or "Unknown", "Other"))
        key = music_key(text)
        index.append({
            "id": eid, "d": "music", "t": f"{style} in {key}" if key else style, "s": snippet(text, f"{style} in {key}" if key else style),
            "f": {"family": family, "style": style, "meter": str(e.get("meter") or "?"),
                  "voices": f"{voices} voice{'s' if voices != 1 else ''}", "tempo": tempo(e.get("bpm")),
                  "language": lang_name(e.get("lang") or lang_of(text))},
        })
        detail[eid] = {
            "brief": text,
            "meta": [["Style", f"{style} ({tag})" if tag and tag != style else style], ["Family", family],
                     ["Key", key or "not stated"], ["Tempo", f"{e.get('bpm')} BPM"], ["Meter", str(e.get("meter"))],
                     ["Voices", str(voices)], ["Length", f"{e.get('length')} bars"], ["Output", "ABC notation"]],
        }
    return index, detail


# ── general ──────────────────────────────────────────────────────────────────
INDUSTRY = {
    "accounting_audit_tax": "Accounting, audit & tax", "finance_insurance": "Finance & insurance",
    "healthcare_ops": "Healthcare operations", "consulting_bizops_analytics": "Consulting & analytics",
    "hr_people": "HR & people", "government_public": "Government & public sector",
    "it_information_systems": "IT & information systems", "education_research_admin": "Education & research admin",
    "energy_utilities_esg": "Energy, utilities & ESG", "legal_compliance": "Legal & compliance",
    "construction_pm": "Construction & project mgmt", "ecommerce_ops": "E-commerce operations",
    "manufacturing_quality": "Manufacturing & quality", "hospitality_fnb": "Hospitality & F&B",
    "last_mile_logistics": "Last-mile logistics", "agriculture_coop": "Agriculture co-ops",
}
ENV_ID = re.compile(r"^s3k_\d+_(?P<slug>.+?)_(?P<lang>en|zh)_(?P<tier>t\d)_rl_\d+$")
FILE_KIND = {
    "xlsx": "Spreadsheet", "xls": "Spreadsheet", "csv": "Spreadsheet", "docx": "Document", "doc": "Document",
    "pdf": "PDF", "pptx": "Slides", "png": "Image", "jpg": "Image", "jpeg": "Image", "svg": "Image",
    "html": "Web page", "xml": "Data", "json": "Data", "zip": "Archive", "md": "Document", "txt": "Document",
}


# display names live in app/names.py so the app and this build agree
sys.path.insert(0, str(Path(__file__).parent))
from app.mimo.names import pretty_system  # noqa: E402


def build_general(df, raw: Path, files: list[str]):
    on_disk = collections.defaultdict(list)
    for f in files:
        m = re.match(r"general/envs/([^/]+)/(.+)", f)
        if m:
            on_disk[m.group(1)].append(m.group(2))

    index, detail = [], {}
    for _, row in df.iterrows():
        j = instance(row)
        text = brief(row)
        eid = j["instance_id"]
        if j.get("dataset_type") == "terminal_bench":
            cat = j.get("category") or "Other"
            sub = j.get("subcategory") or ""
            index.append({
                "id": eid, "d": "general",
                "t": clip(j.get("display_description") or first_sentence(text), 120), "s": snippet(text, clip(j.get("display_description") or first_sentence(text), 120)),
                "f": {"industry": "Terminal tasks", "source": "Terminal-bench style",
                      "tier": "Unrated", "language": lang_name(lang_of(text))},
            })
            detail[eid] = {
                "brief": text,
                "meta": [["Category", f"{cat} / {sub}" if sub else cat], ["Tags", ", ".join(j.get("tags") or [])],
                         ["Resources", f"{j.get('cpus', '?')} CPU · {j.get('memory_mb', '?')} MB · internet {'on' if j.get('allow_internet') else 'off'}"],
                         ["Agent timeout", f"{j.get('agent_timeout_sec', '?')}s"]],
                "files": {"Test files": list(j.get("tests_files") or [])[:30]} if j.get("tests_files") else {},
            }
            continue

        m = ENV_ID.match(eid)
        slug, lang, tier = (m.group("slug"), m.group("lang"), m.group("tier")) if m else ("other", "en", "t?")
        fs = on_disk.get(eid, [])
        systems = sorted({pretty_system(Path(f).stem) for f in fs if f.startswith("tools/") and f.endswith(".py")})
        workspace = sorted(f[len("workspace/"):] for f in fs if f.startswith("workspace/"))
        by_kind = collections.defaultdict(list)
        for w in workspace:
            by_kind[FILE_KIND.get(w.rsplit(".", 1)[-1].lower(), "Other")].append(w)

        rubric = []
        meta_path = raw / "general" / "envs" / eid / "verifier_meta.json"
        if meta_path.exists():
            for it in json.loads(meta_path.read_text()).get("items", []):
                rubric.append({"tier": it.get("tier"), "method": it.get("method"),
                               "weight": it.get("weight"), "q": it.get("question", "")})

        industry = INDUSTRY.get(slug, slug.replace("_", " ").capitalize())
        index.append({
            "id": eid, "d": "general", "t": first_sentence(text, 120), "s": snippet(text, first_sentence(text, 120)),
            "f": {"industry": industry, "source": "Simulated workplace", "tier": f"Tier {tier[1:]}",
                  "language": lang_name(lang)},
            "n": {"systems": len(systems), "files": len(workspace), "checks": len(rubric)},
        })
        detail[eid] = {
            "brief": text,
            "meta": [["Industry", industry], ["Tier", f"Tier {tier[1:]} of 5"], ["Language", lang_name(lang)],
                     ["Grading", f"{len(rubric)} rubric checks"]],
            "systems": systems,
            "files": {k: v for k, v in sorted(by_kind.items(), key=lambda kv: -len(kv[1]))},
            "rubric": rubric,
            "tree": f"https://huggingface.co/datasets/{REPO}/tree/main/general/envs/{eid}",
        }
    return index, detail


# ── reward design ────────────────────────────────────────────────────────────
def build_rewards(raw: Path) -> dict:
    """How General's rubrics are built, across the dataset. Counts only: no question, answer or code."""
    import statistics

    metas = [json.loads(p.read_text()) for p in sorted((raw / "general" / "envs").glob("*/verifier_meta.json"))]
    C = collections.Counter
    per_task, judge_share, files = [], [], C()
    method, tier, mix, inp, act = C(), C(), C(), C(), C()
    weighted = mixed = mixed_heavy = 0
    for m in metas:
        items = m.get("items") or []
        per_task.append(len(items))
        kinds = {i.get("method") for i in items}
        mix["judge_only" if kinds == {"llm"} else "code_only" if kinds == {"rule"} else "both"] += 1
        weighted += any(i.get("weight") is not None for i in items)
        w = lambda i: float(i.get("weight") if i.get("weight") is not None else 1)   # noqa: E731 - verify.py's default
        total = sum(w(i) for i in items) or 1
        # weights written as shares of 1 next to a check with none, which verify.py counts as a full 1
        bare = [i for i in items if i.get("weight") is None]
        if bare and len(bare) < len(items):
            mixed += 1
            mixed_heavy += sum(w(i) for i in bare) / total >= 0.4
        judge_share.append(sum(w(i) for i in items if i.get("method") == "llm") / total)
        inp[m.get("input") or "unknown"] += 1
        act[m.get("act") or "unknown"] += 1
        for i in items:
            method[i.get("method") or "unknown"] += 1
            tier[i.get("tier") or "unknown"] += 1
            if i.get("method") == "llm":
                files.update(Path(f).name for f in i.get("files") or [])
    return {"general": {
        "tasks": len(metas), "checks": sum(per_task),
        "per_task": {"min": min(per_task), "median": statistics.median(per_task), "max": max(per_task)},
        "method": dict(method), "tier": dict(tier), "mix": dict(mix), "weighted_tasks": weighted,
        "mixed_weights": mixed, "mixed_heavy": mixed_heavy,
        "judge_share_median": round(statistics.median(judge_share), 3),
        "input": dict(inp), "act": dict(act), "judged_files": files.most_common(6),
        "uses_gates": sum(any(i.get("gate") for i in m.get("items") or []) or bool(m.get("src_protect")) for m in metas),
    }}


# ── main ─────────────────────────────────────────────────────────────────────
DOMAINS = {
    "code": ("Code", "Fix real issues in real repos", "hidden unit tests", "prog"),
    "webdev": ("Webdev", "Build a website from a brief", "a vision judge", "site"),
    "cyber": ("Cyber", "Reproduce a real crash", "reproducing the crash", "crash"),
    "music": ("Music", "Compose in ABC notation", "a symbolic scorer", "family"),
    "general": ("General", "Knowledge work in a simulated workplace", "a weighted rubric", "industry"),
}
FACET_LABELS = {
    "language": "Brief language", "prog": "Programming language", "crash": "Crash type", "project": "Project", "sanitizer": "Sanitizer",
    "site": "Site type", "framework": "Framework", "style": "Style", "meter": "Meter", "voices": "Voices",
    "tempo": "Tempo", "family": "Family", "industry": "Industry", "source": "Source", "tier": "Tier",
}


# Webdev's tags are read from the brief's wording, not from metadata, so they say so.
LABEL_OVERRIDES = {("webdev", "site"): "Site type (mentioned)", ("webdev", "framework"): "Framework (mentioned)",
                   ("webdev", "style"): "Style (mentioned)"}


def write(path: Path, doc) -> None:
    # mtime=0 so an unchanged build produces byte-identical files (and no spurious commits)
    raw = json.dumps(doc, ensure_ascii=False, separators=(",", ":")).encode()
    path.write_bytes(gzip.compress(raw, compresslevel=9, mtime=0))
    for stale in (path.with_suffix(""),):   # drop the uncompressed copy from earlier builds
        stale.unlink(missing_ok=True)


def main() -> None:
    if "--rewards" in sys.argv:
        from huggingface_hub import snapshot_download

        raw = Path(os.environ.get("RAW_DIR") or snapshot_download(
            REPO, repo_type="dataset", allow_patterns=["general/envs/*/verifier_meta.json"]))
        OUT.mkdir(parents=True, exist_ok=True)
        write(OUT / "rewards.json.gz", build_rewards(raw))
        print(f"rewards  written to {OUT / 'rewards.json.gz'}")
        return
    raw, files = fetch()
    frames = {d: pd.read_parquet(raw / p) for d, p in PARQUETS.items()}
    built = {
        "code": build_code(frames["code"]),
        "webdev": build_webdev(frames["webdev"]),
        "cyber": build_cyber(frames["cyber"]),
        "music": build_music(frames["music"]),
        "general": build_general(frames["general"], raw, files),
    }

    OUT.mkdir(parents=True, exist_ok=True)
    everything, domains = [], []
    for d, (index, detail) in built.items():
        ids = [r["id"] for r in index]
        assert len(ids) == len(set(ids)), f"duplicate ids in {d}"
        everything += index
        write(OUT / f"{d}.json.gz", detail)
        name, task, verifier, main_facet = DOMAINS[d]
        facets = sorted({k for r in index for k in r["f"]}, key=lambda k: (k != main_facet, k))
        domains.append({"id": d, "name": name, "task": task, "verifier": verifier, "count": len(index),
                        "main": main_facet,
                        "facets": [[k, LABEL_OVERRIDES.get((d, k), FACET_LABELS.get(k, k))] for k in facets]})
        print(f"{d:8} {len(index):5} envs  detail {(OUT / f'{d}.json.gz').stat().st_size / 1e6:.2f} MB gz")

    index_doc = {"source": REPO, "total": len(everything), "domains": domains, "envs": everything}
    write(OUT / "index.json.gz", index_doc)
    write(OUT / "rewards.json.gz", build_rewards(raw))
    print(f"index    {len(everything):5} envs  {(OUT / 'index.json.gz').stat().st_size / 1e6:.2f} MB gz")


if __name__ == "__main__":
    main()