#!/usr/bin/env python3
"""MLX Model Explorer community benchmark.

Measures one MLX model on this Mac with mlx-lm and prints an anonymous JSON result:
prompt speed, generation speed, time to first token and peak memory. It uses the
same approach as `mlx_lm.benchmark` (random prompt tokens, EOS disabled, warm-up
run, several timed trials) and reports medians.

Nothing is sent anywhere unless you pass --submit URL, and even then the exact
payload is shown and you are asked to confirm. The payload holds no hostname,
username, paths or serial numbers: only the chip name (e.g. "Apple M3 Max"), the
RAM size class, OS major version and library versions.

    pip install -U mlx-lm
    python mlx_explorer_bench.py --model mlx-community/Qwen3-0.6B-4bit
    python mlx_explorer_bench.py --model mlx-community/Qwen3-0.6B-4bit --submit https://<space-url>
"""

from __future__ import annotations

import argparse
import json
import platform
import statistics
import subprocess
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

BENCHMARK_VERSION = "mlxbench-2"  # 2: every perplexity window starts with BOS

# Quality: perplexity on a fixed slice of the wikitext-2 test set (CC BY-SA 3.0).
# Perplexity depends on the tokenizer, so results are only compared between
# quantizations of the same base model.
EVAL_DATASET = "wikitext2-test-v1"
EVAL_FILE = "eval/wikitext-2-raw-v1-test-v1.txt"
EVAL_SHA256 = "5dacf0e5f3a9a8be86bbf99aae316fb2849f84fc62b913ab3b378d7f03c730e8"
EVAL_REPO = "mlx-community/mlx-model-explorer-data"
RAM_CLASSES = [8, 16, 18, 24, 32, 36, 48, 64, 96, 128, 192, 256, 512]
CONTEXTS = [4096, 8192, 16384, 32768, 65536, 131072, 262144]


def sysctl(name: str) -> str | None:
    try:
        return subprocess.run(["sysctl", "-n", name], capture_output=True, text=True, timeout=5).stdout.strip() or None
    except (OSError, subprocess.SubprocessError):
        return None


def chip_name() -> str:
    brand = sysctl("machdep.cpu.brand_string") or ""
    import re

    m = re.match(r"^Apple (M[1-9])( (Pro|Max|Ultra))?", brand)
    return f"Apple {m.group(1)}{m.group(2) or ''}" if m else "other"


def ram_class() -> int | None:
    raw = sysctl("hw.memsize")
    if not raw or not raw.isdigit():
        return None
    gb = int(raw) / 1024**3
    return min(RAM_CLASSES, key=lambda c: abs(c - gb))


def context_bucket(tokens: int) -> int:
    return next((c for c in CONTEXTS if c >= tokens), CONTEXTS[-1])


def load_eval_text() -> str:
    import hashlib

    local = Path(__file__).resolve().parent / EVAL_FILE
    if local.exists():
        data = local.read_bytes()
    else:
        from huggingface_hub import hf_hub_download

        data = Path(hf_hub_download(EVAL_REPO, EVAL_FILE, repo_type="dataset")).read_bytes()
    if hashlib.sha256(data).hexdigest() != EVAL_SHA256:
        raise SystemExit("evaluation text does not match the expected checksum; refusing to measure")
    return data.decode()


def perplexity(model, tokenizer, max_tokens: int, window: int) -> tuple[float, float, int]:
    """Mean next-token cross-entropy over non-overlapping windows. Returns (ppl, stderr, tokens scored)."""
    import math

    import mlx.core as mx
    import mlx.nn as nn

    ids = tokenizer.encode(load_eval_text())
    # Models with a BOS token (Gemma, LFM, MiniCPM, ...) expect it at the start of
    # every sequence; without it Gemma's perplexity runs into the thousands.
    # Some tokenizers add it on encode, some don't, so strip it and add it per window.
    bos = getattr(tokenizer, "bos_token_id", None)
    if bos is not None and ids and ids[0] == bos:
        ids = ids[1:]
    step = window - 1 if bos is not None else window
    n_windows = min(len(ids), max_tokens) // step
    if n_windows < 1:
        raise SystemExit("evaluation text is shorter than one window")
    losses = []
    for w in range(n_windows):
        chunk = mx.array(([bos] if bos is not None else []) + ids[w * step:(w + 1) * step])[None]
        logits = model(chunk[:, :-1]).astype(mx.float32)
        loss = nn.losses.cross_entropy(logits, chunk[:, 1:], reduction="none")
        mx.eval(loss)
        losses.append(loss.flatten())
        print(f"  perplexity window {w + 1}/{n_windows}", file=sys.stderr, end="\r")
    print(file=sys.stderr)
    all_losses = mx.concatenate(losses)
    mean = all_losses.mean().item()
    std = mx.sqrt(mx.var(all_losses)).item()
    count = all_losses.size
    ppl = math.exp(mean)
    return ppl, ppl * std / math.sqrt(count), count


def run(args) -> dict:
    import mlx.core as mx
    import mlx_lm
    from mlx_lm import load, stream_generate

    mx.random.seed(0)
    print(f"Loading {args.model} …", file=sys.stderr)
    model, tokenizer, config = load(args.model, return_config=True)
    tokenizer._eos_token_ids = {}  # never stop early, like mlx_lm.benchmark
    vocab = config.get("vocab_size") or (config.get("text_config") or {}).get("vocab_size")
    prompt = mx.random.randint(0, vocab, (args.prompt_tokens,)).tolist()

    def once():
        start = time.perf_counter()
        first = None
        last = None
        for resp in stream_generate(model, tokenizer, prompt, max_tokens=args.generation_tokens):
            if first is None:
                first = time.perf_counter() - start
            last = resp
        return last, first

    result = {}
    if args.quality:
        print(f"Measuring perplexity on {EVAL_DATASET} ({args.ppl_tokens} tokens) …", file=sys.stderr)
        ppl, err, scored = perplexity(model, tokenizer, args.ppl_tokens, args.ppl_window)
        print(f"perplexity {ppl:.3f} ± {err:.3f} over {scored} tokens", file=sys.stderr)
        result.update({"perplexity": round(ppl, 4), "perplexity_stderr": round(err, 4),
                       "eval_dataset": EVAL_DATASET, "eval_tokens": int(scored)})

    print("Warm-up run …", file=sys.stderr)
    once()
    mx.reset_peak_memory()
    trials = []
    for i in range(args.trials):
        resp, ttft = once()
        trials.append((resp, ttft))
        print(f"trial {i + 1}: prompt {resp.prompt_tps:.1f} tok/s, generation {resp.generation_tps:.1f} tok/s, "
              f"TTFT {ttft * 1000:.0f} ms, peak {resp.peak_memory:.2f} GB", file=sys.stderr)

    med = lambda xs: statistics.median(xs)
    mac = platform.mac_ver()[0]
    return result | {
        "event_type": "mlx_benchmark",
        "benchmark_type": "mlx_lm",
        "benchmark_version": BENCHMARK_VERSION,
        "selected_model": args.model,
        "prompt_tokens": args.prompt_tokens,
        "generation_tokens": args.generation_tokens,
        "prompt_tps": round(med([r.prompt_tps for r, _ in trials]), 2),
        "generation_tps": round(med([r.generation_tps for r, _ in trials]), 2),
        "ttft_ms": round(med([t for _, t in trials]) * 1000, 1),
        "peak_memory_gb": round(max(r.peak_memory for r, _ in trials), 3),
        "target_context": context_bucket(args.prompt_tokens + args.generation_tokens),
        "chip": chip_name(),
        "reported_ram_gb": ram_class(),
        "macos_major": int(mac.split(".")[0]) if mac and mac.split(".")[0].isdigit() else None,
        "mlx_version": getattr(mx, "__version__", None),
        "mlx_lm_version": getattr(mlx_lm, "__version__", None),
    }


def submit(result: dict, url: str, token: str | None, assume_yes: bool) -> int:
    payload = {"events": [{k: v for k, v in result.items() if v is not None}]}
    body = json.dumps(payload).encode()
    print("\nThis exact payload will be sent to", url.rstrip("/") + "/api/events", file=sys.stderr)
    print(json.dumps(payload, indent=2), file=sys.stderr)
    if not assume_yes:
        if input("Send it? [y/N] ").strip().lower() not in ("y", "yes"):
            print("Not sent.", file=sys.stderr)
            return 0
    headers = {"content-type": "application/json", "user-agent": f"mlx-explorer-bench/{BENCHMARK_VERSION}"}
    if token:
        headers["authorization"] = f"Bearer {token}"
    req = urllib.request.Request(url.rstrip("/") + "/api/events", data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(req, timeout=30) as r:
            reply = json.loads(r.read() or b"{}")
            flags = (reply.get("flags") or [[]])[0]
            print("Submitted." + (f" Flagged for review: {', '.join(flags)}" if flags else " Thank you!"), file=sys.stderr)
            return 0
    except urllib.error.HTTPError as e:
        print(f"Not accepted: HTTP {e.code} {e.read()[:500].decode(errors='replace')}", file=sys.stderr)
        return 1
    except urllib.error.URLError as e:
        print(f"Couldn't reach {url}: {e.reason}", file=sys.stderr)
        return 1


def main(argv=None) -> int:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument("--model", required=True, help="Hugging Face repo id, e.g. mlx-community/Qwen3-0.6B-4bit")
    p.add_argument("--prompt-tokens", "-p", type=int, default=512)
    p.add_argument("--generation-tokens", "-g", type=int, default=256)
    p.add_argument("--trials", "-n", type=int, default=3)
    p.add_argument("--quality", action="store_true",
                   help=f"also measure perplexity on {EVAL_DATASET} (compare only quants of the same base model)")
    p.add_argument("--ppl-tokens", type=int, default=16384, help="tokens of evaluation text to score")
    p.add_argument("--ppl-window", type=int, default=1024, help="window length for perplexity")
    p.add_argument("--submit", metavar="URL", help="MLX Model Explorer URL to submit the result to")
    p.add_argument("--yes", action="store_true", help="submit without asking for confirmation")
    p.add_argument("--hf-token", default=None,
                   help="only needed while the Space is private: a Hugging Face token sent as a bearer header")
    args = p.parse_args(argv)
    if not (1 <= args.prompt_tokens <= 1_048_576 and 1 <= args.generation_tokens <= 65_536 and 1 <= args.trials <= 20
            and 1024 <= args.ppl_tokens <= 262_144 and 128 <= args.ppl_window <= 8192):
        p.error("prompt/generation tokens or trials out of range")
    if platform.system() != "Darwin" or platform.machine() != "arm64":
        print("MLX runs on Apple Silicon Macs; this machine can't run the benchmark.", file=sys.stderr)
        return 2
    result = run(args)
    print(json.dumps(result, indent=2))
    if args.submit:
        return submit(result, args.submit, args.hf_token, args.yes)
    print("\nTo contribute this result, paste the JSON into MLX Model Explorer or rerun with --submit URL.", file=sys.stderr)
    return 0


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