"""
Comprehensive benchmark: nanotok vs tiktoken vs HuggingFace tokenizers vs transformers.

Tests correctness (parity) and throughput across many model families:
  LLaMA, Qwen, DeepSeek, GLM, GPT (tiktoken), Arcee, Mistral, Gemma, Phi, etc.

Usage:
  python benchmarks/bench.py
  python benchmarks/bench.py --quick          # fewer iterations
  python benchmarks/bench.py --models llama   # specific family
"""

from __future__ import annotations

import argparse
import gc
import random
import string
import sys
import time
from dataclasses import dataclass, field

# ── Test corpus ──────────────────────────────────────────────────────────────

UNICODE_SAMPLES = [
    "こんにちは世界 Hello World 你好世界",
    "Привет мир! مرحبا بالعالم 🌍🚀",
    "👨‍👩‍👧‍👦 family emoji test 🏳️‍🌈",
    "€1,234.56 £999 ¥10000 ₹500",
    "안녕하세요 세계! สวัสดีชาวโลก",
    "αβγδ εζηθ ικλμ νξοπ",
    "नमस्ते दुनिया שלום עולם",
    "𝕳𝖊𝖑𝖑𝖔 𝕿𝖊𝖘𝖙 𝕴𝖓𝖕𝖚𝖙",
    "\t\ttabs\tand\nnewlines\r\nand\r carriage returns",
    "mixed: café résumé naïve Ñoño über straße",
]

CODE_SAMPLES = [
    'def fibonacci(n):\n    a, b = 0, 1\n    for _ in range(n):\n        a, b = b, a + b\n    return a\n\nprint(fibonacci(100))',
    '#include <iostream>\nint main() {\n    std::cout << "Hello, World!" << std::endl;\n    return 0;\n}\n',
    'SELECT u.name, COUNT(o.id) AS order_count\nFROM users u\nLEFT JOIN orders o ON u.id = o.user_id\nWHERE u.created_at > \'2024-01-01\'\nGROUP BY u.name\nHAVING COUNT(o.id) > 5\nORDER BY order_count DESC;',
    '{\n  "model": "gpt-4",\n  "messages": [\n    {"role": "system", "content": "You are helpful."},\n    {"role": "user", "content": "Hello!"}\n  ],\n  "temperature": 0.7\n}',
]

LONG_TEXT = (
    "The quick brown fox jumps over the lazy dog. " * 200
    + "Pack my box with five dozen liquor jugs. " * 100
    + "How vexingly quick daft zebras jump! " * 150
)


def build_corpus(rng: random.Random, n_random: int = 300) -> list[str]:
    alpha = string.ascii_letters + string.digits + string.punctuation + " \t\n"
    corpus = list(UNICODE_SAMPLES) + list(CODE_SAMPLES) + [LONG_TEXT]

    for _ in range(n_random):
        length = rng.randint(1, 2000)
        text = "".join(rng.choice(alpha) for _ in range(length))
        if rng.random() < 0.3:
            text += rng.choice(UNICODE_SAMPLES)
        corpus.append(text)

    corpus.append("")
    corpus.append("a")
    corpus.append(" ")
    corpus.append("   \n\n\t  ")
    corpus.append("a" * 10000)
    corpus.append("ABCD" * 3000)
    corpus.append("\n".join(f"line {i}: content" for i in range(500)))

    return corpus


# ── Results ──────────────────────────────────────────────────────────────────

@dataclass
class BenchResult:
    name: str
    backend: str
    parity: bool = True
    mismatches: int = 0
    total: int = 0
    nanotok_enc_ms: float = 0.0
    baseline_enc_ms: float = 0.0
    nanotok_dec_ms: float = 0.0
    baseline_dec_ms: float = 0.0
    speedup_enc: float = 0.0
    speedup_dec: float = 0.0
    errors: list[str] = field(default_factory=list)

    @property
    def parity_pct(self) -> float:
        return ((self.total - self.mismatches) / self.total * 100) if self.total else 0


# ── Timing helper ────────────────────────────────────────────────────────────

def _time_encode(encode_fn, texts: list[str], warmup: int = 2, runs: int = 5) -> float:
    for _ in range(warmup):
        for t in texts[:50]:
            encode_fn(t)

    gc.disable()
    best = float("inf")
    for _ in range(runs):
        start = time.perf_counter()
        for t in texts:
            encode_fn(t)
        elapsed = (time.perf_counter() - start) * 1000
        best = min(best, elapsed)
    gc.enable()
    return best


def _time_decode(decode_fn, all_ids: list[list[int]], warmup: int = 2, runs: int = 5) -> float:
    for _ in range(warmup):
        for ids in all_ids[:50]:
            decode_fn(ids)

    gc.disable()
    best = float("inf")
    for _ in range(runs):
        start = time.perf_counter()
        for ids in all_ids:
            decode_fn(ids)
        elapsed = (time.perf_counter() - start) * 1000
        best = min(best, elapsed)
    gc.enable()
    return best


# ── HuggingFace model benchmarks ────────────────────────────────────────────

HF_MODELS = {
    "llama": [
        ("LLaMA-3.1-8B-Inst", "meta-llama/Llama-3.1-8B-Instruct"),
    ],
    "qwen": [
        ("Qwen2.5-0.5B-Inst", "Qwen/Qwen2.5-0.5B-Instruct"),
    ],
    "deepseek": [
        ("DeepSeek-R1-0528", "deepseek-ai/DeepSeek-R1-0528"),
    ],
    "smollm": [
        ("SmolLM2-135M-Inst", "HuggingFaceTB/SmolLM2-135M-Instruct"),
    ],
    "arcee": [
        ("Arcee-Nova", "arcee-ai/Arcee-Nova"),
    ],
    "mistral": [
        ("Mistral-7B-Inst-v0.3", "mistralai/Mistral-7B-Instruct-v0.3"),
    ],
    "gemma": [
        ("Gemma-2-2B-IT", "google/gemma-2-2b-it"),
    ],
    "phi": [
        ("Phi-3-mini-4k", "microsoft/Phi-3-mini-4k-instruct"),
    ],
    "gpt_oss": [
        ("GPT-NeoX-20B", "EleutherAI/gpt-neox-20b"),
    ],
}


def bench_hf_model(
    label: str, repo: str, corpus: list[str], *, time_it: bool = True
) -> BenchResult:
    from tokenizers import Tokenizer as HFTokenizer
    from nanotok import Tokenizer

    result = BenchResult(name=label, backend="HF tokenizers")

    try:
        tok_hf = HFTokenizer.from_pretrained(repo)
    except Exception as e:
        result.errors.append(f"HF load failed: {e}")
        return result

    try:
        tok_nt = Tokenizer.from_pretrained(repo)
    except Exception as e:
        result.errors.append(f"nanotok load failed: {e}")
        return result

    result.total = len(corpus)
    first_mismatch_shown = False
    for text in corpus:
        try:
            ref_ids = tok_hf.encode(text, add_special_tokens=False).ids
            got_ids = tok_nt.encode(text, add_special_tokens=False)
        except Exception as e:
            result.mismatches += 1
            if not first_mismatch_shown:
                result.errors.append(f"encode exception: {e}")
                first_mismatch_shown = True
            continue

        if ref_ids != got_ids:
            result.mismatches += 1
            if not first_mismatch_shown:
                snippet = text[:80].replace("\n", "\\n")
                result.errors.append(
                    f"first mismatch: text={snippet!r} ref_len={len(ref_ids)} got_len={len(got_ids)}"
                )
                first_mismatch_shown = True

        try:
            ref_decoded = tok_hf.decode(ref_ids, skip_special_tokens=False)
            got_decoded = tok_nt.decode(ref_ids, skip_special_tokens=False)
            if ref_decoded != got_decoded:
                if len(result.errors) < 3:
                    snippet = text[:60].replace("\n", "\\n")
                    result.errors.append(f"decode mismatch: text={snippet!r}")
        except Exception:
            pass

    result.parity = result.mismatches == 0

    if time_it and result.mismatches < result.total:
        try:
            result.nanotok_enc_ms = _time_encode(
                lambda t: tok_nt.encode(t, add_special_tokens=False), corpus
            )
            result.baseline_enc_ms = _time_encode(
                lambda t: tok_hf.encode(t, add_special_tokens=False), corpus
            )
            result.speedup_enc = result.baseline_enc_ms / result.nanotok_enc_ms if result.nanotok_enc_ms > 0 else 0

            all_ids = [tok_nt.encode(t, add_special_tokens=False) for t in corpus]
            result.nanotok_dec_ms = _time_decode(lambda ids: tok_nt.decode(ids), all_ids)
            result.baseline_dec_ms = _time_decode(lambda ids: tok_hf.decode(ids), all_ids)
            result.speedup_dec = result.baseline_dec_ms / result.nanotok_dec_ms if result.nanotok_dec_ms > 0 else 0
        except Exception as e:
            result.errors.append(f"timing failed: {e}")

    return result


# ── Tiktoken benchmarks ─────────────────────────────────────────────────────

TIKTOKEN_ENCODINGS = {
    "gpt": [
        ("cl100k_base (GPT-4)", "cl100k_base"),
        ("o200k_base (GPT-4o)", "o200k_base"),
        ("p50k_base (GPT-3.5)", "p50k_base"),
        ("r50k_base (GPT-3)", "r50k_base"),
        ("gpt2", "gpt2"),
    ],
}


def bench_tiktoken(
    label: str, encoding: str, corpus: list[str], *, time_it: bool = True
) -> BenchResult:
    import tiktoken
    from nanotok import Tokenizer

    result = BenchResult(name=label, backend="tiktoken")

    try:
        tok_tt = tiktoken.get_encoding(encoding)
    except Exception as e:
        result.errors.append(f"tiktoken load failed: {e}")
        return result

    try:
        tok_nt = Tokenizer.from_tiktoken(encoding)
    except Exception as e:
        result.errors.append(f"nanotok load failed: {e}")
        return result

    result.total = len(corpus)
    first_mismatch_shown = False
    for text in corpus:
        try:
            ref_ids = tok_tt.encode(text)
            got_ids = tok_nt.encode(text)
        except Exception as e:
            result.mismatches += 1
            if not first_mismatch_shown:
                result.errors.append(f"encode exception: {e}")
                first_mismatch_shown = True
            continue

        if ref_ids != got_ids:
            result.mismatches += 1
            if not first_mismatch_shown:
                snippet = text[:80].replace("\n", "\\n")
                result.errors.append(
                    f"first mismatch: text={snippet!r} ref_len={len(ref_ids)} got_len={len(got_ids)}"
                )
                first_mismatch_shown = True

    result.parity = result.mismatches == 0

    if time_it and result.mismatches < result.total:
        try:
            result.nanotok_enc_ms = _time_encode(lambda t: tok_nt.encode(t), corpus)
            result.baseline_enc_ms = _time_encode(lambda t: tok_tt.encode(t), corpus)
            result.speedup_enc = result.baseline_enc_ms / result.nanotok_enc_ms if result.nanotok_enc_ms > 0 else 0

            all_ids = [tok_nt.encode(t) for t in corpus]
            result.nanotok_dec_ms = _time_decode(lambda ids: tok_nt.decode(ids), all_ids)
            result.baseline_dec_ms = _time_decode(lambda ids: tok_tt.decode(ids), all_ids)
            result.speedup_dec = result.baseline_dec_ms / result.nanotok_dec_ms if result.nanotok_dec_ms > 0 else 0
        except Exception as e:
            result.errors.append(f"timing failed: {e}")

    return result


# ── Transformers benchmark (AutoTokenizer) ───────────────────────────────────

TRANSFORMERS_MODELS = {
    "llama": ("LLaMA-3.1-8B (transformers)", "meta-llama/Llama-3.1-8B-Instruct"),
    "qwen": ("Qwen2.5-0.5B (transformers)", "Qwen/Qwen2.5-0.5B-Instruct"),
    "deepseek": ("DeepSeek-R1-0528 (transformers)", "deepseek-ai/DeepSeek-R1-0528"),
    "mistral": ("Mistral-7B-v0.3 (transformers)", "mistralai/Mistral-7B-Instruct-v0.3"),
}


def bench_transformers(
    label: str, repo: str, corpus: list[str], *, time_it: bool = True
) -> BenchResult:
    from transformers import AutoTokenizer
    from nanotok import Tokenizer

    result = BenchResult(name=label, backend="transformers")

    try:
        tok_tf = AutoTokenizer.from_pretrained(repo)
    except Exception as e:
        result.errors.append(f"transformers load failed: {e}")
        return result

    try:
        tok_nt = Tokenizer.from_pretrained(repo)
    except Exception as e:
        result.errors.append(f"nanotok load failed: {e}")
        return result

    result.total = len(corpus)
    first_mismatch_shown = False
    for text in corpus:
        try:
            ref_ids = tok_tf.encode(text, add_special_tokens=False)
            got_ids = tok_nt.encode(text, add_special_tokens=False)
        except Exception as e:
            result.mismatches += 1
            if not first_mismatch_shown:
                result.errors.append(f"encode exception: {e}")
                first_mismatch_shown = True
            continue

        if ref_ids != got_ids:
            result.mismatches += 1
            if not first_mismatch_shown:
                snippet = text[:80].replace("\n", "\\n")
                result.errors.append(
                    f"first mismatch: text={snippet!r} ref_len={len(ref_ids)} got_len={len(got_ids)}"
                )
                first_mismatch_shown = True

    result.parity = result.mismatches == 0

    if time_it and result.mismatches < result.total:
        try:
            result.nanotok_enc_ms = _time_encode(
                lambda t: tok_nt.encode(t, add_special_tokens=False), corpus
            )
            result.baseline_enc_ms = _time_encode(
                lambda t: tok_tf.encode(t, add_special_tokens=False), corpus
            )
            result.speedup_enc = result.baseline_enc_ms / result.nanotok_enc_ms if result.nanotok_enc_ms > 0 else 0

            all_ids = [tok_nt.encode(t, add_special_tokens=False) for t in corpus]
            result.nanotok_dec_ms = _time_decode(lambda ids: tok_nt.decode(ids), all_ids)
            result.baseline_dec_ms = _time_decode(lambda ids: tok_tf.decode(ids), all_ids)
            result.speedup_dec = result.baseline_dec_ms / result.nanotok_dec_ms if result.nanotok_dec_ms > 0 else 0
        except Exception as e:
            result.errors.append(f"timing failed: {e}")

    return result


# ── Display ──────────────────────────────────────────────────────────────────

def print_header(title: str):
    print(f"\n{'='*80}")
    print(f"  {title}")
    print(f"{'='*80}")


def print_result(r: BenchResult):
    status = "PASS" if r.parity else "FAIL"
    parity_str = f"{r.parity_pct:.1f}% ({r.total - r.mismatches}/{r.total})"

    print(f"\n  {r.name}  vs  {r.backend}")
    print(f"  {'─'*60}")
    print(f"  Parity:        [{status}] {parity_str}")

    if r.nanotok_enc_ms > 0:
        print(f"  Encode:        nanotok {r.nanotok_enc_ms:>8.1f} ms | {r.backend:>14s} {r.baseline_enc_ms:>8.1f} ms | {r.speedup_enc:>5.1f}x")
    if r.nanotok_dec_ms > 0:
        print(f"  Decode:        nanotok {r.nanotok_dec_ms:>8.1f} ms | {r.backend:>14s} {r.baseline_dec_ms:>8.1f} ms | {r.speedup_dec:>5.1f}x")

    for err in r.errors:
        print(f"  >> {err}")


def print_summary(results: list[BenchResult]):
    print_header("SUMMARY")

    passed = sum(1 for r in results if r.parity)
    failed = sum(1 for r in results if not r.parity and r.total > 0)
    errored = sum(1 for r in results if r.total == 0 and r.errors)

    print(f"\n  Total: {len(results)}  |  Passed: {passed}  |  Failed: {failed}  |  Errored: {errored}")

    enc_speedups = [r.speedup_enc for r in results if r.speedup_enc > 0]
    dec_speedups = [r.speedup_dec for r in results if r.speedup_dec > 0]

    if enc_speedups:
        avg_enc = sum(enc_speedups) / len(enc_speedups)
        min_enc = min(enc_speedups)
        max_enc = max(enc_speedups)
        print(f"\n  Encode speedup:  avg {avg_enc:.1f}x  |  min {min_enc:.1f}x  |  max {max_enc:.1f}x")

    if dec_speedups:
        avg_dec = sum(dec_speedups) / len(dec_speedups)
        min_dec = min(dec_speedups)
        max_dec = max(dec_speedups)
        print(f"  Decode speedup:  avg {avg_dec:.1f}x  |  min {min_dec:.1f}x  |  max {max_dec:.1f}x")

    print(f"\n{'─'*80}")
    fmt = "  {:<35s} {:>7s} {:>10s} {:>10s} {:>8s}"
    print(fmt.format("Model", "Parity", "Enc (ms)", "Base (ms)", "Speedup"))
    print(f"  {'─'*72}")
    for r in results:
        if r.total == 0 and r.errors:
            print(f"  {r.name:<35s} {'ERROR':>7s}   {r.errors[0][:40]}")
        else:
            status = "PASS" if r.parity else "FAIL"
            enc_str = f"{r.nanotok_enc_ms:.1f}" if r.nanotok_enc_ms > 0 else "-"
            base_str = f"{r.baseline_enc_ms:.1f}" if r.baseline_enc_ms > 0 else "-"
            spd_str = f"{r.speedup_enc:.1f}x" if r.speedup_enc > 0 else "-"
            print(fmt.format(r.name, status, enc_str, base_str, spd_str))

    print()


# ── Main ─────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="nanotok benchmark suite")
    ap.add_argument("--models", default="all", help="Comma-separated families: llama,qwen,deepseek,glm,gpt,gpt_oss,arcee,mistral,gemma,phi (or 'all')")
    ap.add_argument("--quick", action="store_true", help="Fewer random samples, skip timing")
    ap.add_argument("--no-timing", action="store_true", help="Skip throughput benchmarks")
    ap.add_argument("--no-transformers", action="store_true", help="Skip transformers comparison")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    rng = random.Random(args.seed)
    n_random = 50 if args.quick else 300
    corpus = build_corpus(rng, n_random=n_random)
    time_it = not args.quick and not args.no_timing

    families = [f.strip() for f in args.models.split(",")]
    if "all" in families:
        families = list(HF_MODELS.keys()) + ["gpt"]

    results: list[BenchResult] = []

    # HF tokenizers comparison
    for family in families:
        if family == "gpt":
            continue
        if family not in HF_MODELS:
            print(f"  [skip] unknown family: {family}")
            continue

        print_header(f"HuggingFace tokenizers: {family.upper()}")
        for label, repo in HF_MODELS[family]:
            print(f"\n  Loading {repo} ...")
            r = bench_hf_model(label, repo, corpus, time_it=time_it)
            print_result(r)
            results.append(r)

    # tiktoken comparison
    if "gpt" in families:
        print_header("tiktoken (OpenAI GPT)")
        for label, encoding in TIKTOKEN_ENCODINGS["gpt"]:
            print(f"\n  Loading {encoding} ...")
            r = bench_tiktoken(label, encoding, corpus, time_it=time_it)
            print_result(r)
            results.append(r)

    # transformers comparison (subset)
    if not args.no_transformers:
        active_tf = {k: v for k, v in TRANSFORMERS_MODELS.items() if k in families}
        if active_tf:
            print_header("transformers AutoTokenizer")
            for family_key, (label, repo) in active_tf.items():
                print(f"\n  Loading {repo} ...")
                r = bench_transformers(label, repo, corpus, time_it=time_it)
                print_result(r)
                results.append(r)

    print_summary(results)

    any_fail = any(not r.parity for r in results if r.total > 0)
    sys.exit(1 if any_fail else 0)


if __name__ == "__main__":
    main()
