"""
Stress + parity testing for nanotok against strong baselines.

Usage:
  python benchmarks/stress.py --models Qwen/Qwen2.5-0.5B-Instruct,meta-llama/Llama-3.1-8B-Instruct
  python benchmarks/stress.py --encoding cl100k_base
"""

from __future__ import annotations

import argparse
import os
import random
import string
from dataclasses import dataclass
from pathlib import Path


def _random_unicode_pool() -> list[str]:
    return [
        "こんにちは世界",
        "你好世界",
        "안녕하세요",
        "Привет мир",
        "مرحبا بالعالم",
        "नमस्ते दुनिया",
        "שלום עולם",
        "γειά σου κόσμε",
        "สวัสดีชาวโลก",
        "😀😃😄😁😆😅😂🤣😊😇",
        "👨\u200d👩\u200d👧\u200d👦",
        "🏳️\u200d🌈",
        "€£¥₩₽₹",
    ]


def _adversarial_cases() -> list[str]:
    return [
        "a" * 10000,
        "A" * 10000,
        "_" * 10000,
        "0" * 10000,
        "hello" * 2000,
        "The_quick_brown_fox" * 1000,
        "ABCD" * 4000,
        ("\n" * 1000) + "end",
        (" \t\r\n" * 2000),
        "".join(random.choice(string.ascii_letters + string.digits) for _ in range(5000)),
        "".join(random.choice(string.ascii_letters + string.digits + "+/") for _ in range(8000)),
    ]


def _random_text(rng: random.Random, min_len: int, max_len: int) -> str:
    length = rng.randint(min_len, max_len)
    alpha = string.ascii_letters + string.digits + string.punctuation + " \t\n\r"
    base = "".join(rng.choice(alpha) for _ in range(length))
    if rng.random() < 0.3:
        base += rng.choice(_random_unicode_pool())
    if rng.random() < 0.2:
        base = rng.choice(_random_unicode_pool()) + base
    return base


@dataclass
class ParityResult:
    ok: bool
    failures: int
    total: int


def _load_hf_tokenizer(model: str):
    from tokenizers import Tokenizer as HFTokenizer

    tok = HFTokenizer.from_pretrained(model)
    return tok


def _load_nanotok(model: str):
    from nanotok import Tokenizer

    return Tokenizer.from_pretrained(model)


def _load_nanotok_tiktoken(encoding: str):
    from nanotok import Tokenizer

    return Tokenizer.from_tiktoken(encoding)


def _load_tiktoken(encoding: str):
    import tiktoken

    return tiktoken.get_encoding(encoding)


def _chunks_consistency(tok, text: str) -> None:
    # Internal consistency check: encode == encode_chunks(debug_chunks)
    chunks = tok._engine.debug_chunks(text)
    ids_a = tok.encode(text)
    ids_b = tok._engine.encode_chunks(chunks)
    if ids_a != ids_b:
        raise AssertionError("encode != encode_chunks(debug_chunks)")


def _parity_hf(model: str, texts: list[str], *, check_chunks: bool) -> ParityResult:
    tok_ref = _load_hf_tokenizer(model)
    tok = _load_nanotok(model)
    failures = 0
    for t in texts:
        ref = tok_ref.encode(t, add_special_tokens=False).ids
        got = tok.encode(t, add_special_tokens=False)
        if ref != got:
            failures += 1
        if check_chunks:
            _chunks_consistency(tok, t)
    return ParityResult(failures == 0, failures, len(texts))


def _parity_tiktoken(encoding: str, texts: list[str]) -> ParityResult:
    ref = _load_tiktoken(encoding)
    tok = _load_nanotok_tiktoken(encoding)
    failures = 0
    for t in texts:
        ref_ids = ref.encode(t)
        got = tok.encode(t)
        if ref_ids != got:
            failures += 1
    return ParityResult(failures == 0, failures, len(texts))


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", default="")
    ap.add_argument("--encoding", default="")
    ap.add_argument("--iters", type=int, default=200)
    ap.add_argument("--min-len", type=int, default=1)
    ap.add_argument("--max-len", type=int, default=512)
    ap.add_argument("--seed", type=int, default=1337)
    ap.add_argument("--include-adversarial", action="store_true")
    ap.add_argument("--check-chunks", action="store_true")
    args = ap.parse_args()

    rng = random.Random(args.seed)
    texts = [_random_text(rng, args.min_len, args.max_len) for _ in range(args.iters)]
    if args.include_adversarial:
        texts += _adversarial_cases()

    if args.models:
        for model in args.models.split(","):
            model = model.strip()
            if not model:
                continue
            print(f"HF parity: {model}")
            r = _parity_hf(model, texts, check_chunks=args.check_chunks)
            print(f"  ok={r.ok} failures={r.failures}/{r.total}")

    if args.encoding:
        print(f"tiktoken parity: {args.encoding}")
        r = _parity_tiktoken(args.encoding, texts)
        print(f"  ok={r.ok} failures={r.failures}/{r.total}")


if __name__ == "__main__":
    main()
