#!/usr/bin/env python3
"""Compare the Agentic 32K tokenizer with common tiktoken encodings."""

from __future__ import annotations

import argparse
import json
import unicodedata
from collections import defaultdict
from pathlib import Path

import tiktoken
from tokenizers import Tokenizer

SUITE: dict[str, tuple[str, ...]] = {
    "english": (
        "A reliable agent should inspect the environment, make a concise plan, execute only "
        "authorized actions, verify the observed result, and report limitations honestly.",
        "Deterministic multi-hash routing assigns every token to a reproducible pair of experts "
        "without a learned router or an auxiliary load-balancing loss.",
    ),
    "french": (
        "Un agent fiable doit comprendre la demande, vérifier les contraintes, exécuter les "
        "étapes autorisées puis confirmer le résultat avec des preuves reproductibles.",
        "Le modèle compact conserve une architecture inspectable, un routage déterministe et "
        "une chaîne de publication reliant les poids, le code, les données et les évaluations.",
    ),
    "code": (
        "async def fetch_json(session, url):\n"
        "    async with session.get(url, timeout=30) as response:\n"
        "        response.raise_for_status()\n"
        "        return await response.json()\n",
        "def verify_shard(path, expected_sha256):\n"
        "    digest = hashlib.sha256(path.read_bytes()).hexdigest()\n"
        "    assert digest == expected_sha256, (path, digest)\n",
    ),
    "json_tool": (
        '{"name":"search","arguments":{"query":"TR-HASH tokenizer efficiency",'
        '"top_k":5},"request_id":"req_1729"}',
        '{"status":"ok","result":{"verified":true,"rows":125000000000,'
        '"sha256":"d2053a9e99c484f7c"}}',
    ),
    "math": (
        "Solve 3x² - 14x - 5 = 0. The discriminant is Δ = (-14)² - 4·3·(-5) = 256, "
        "so x = (14 ± 16)/6 and the solutions are 5 and -1/3.",
        "For every ε > 0 there exists δ > 0 such that |x-a| < δ implies |f(x)-f(a)| < ε.",
    ),
    "emoji": (
        "😀 😃 😄 😁 😂 🥹 😊 🚀 ✅ ❌ ⚠️ 🔧 🧠 🤖 📊 🇫🇷",
        "Bonjour 👋🏽 — développeuse 🧑🏽‍💻, famille 👨‍👩‍👧‍👦, science 🧬🔬 et feu ❤️‍🔥.",
    ),
    "agentic_markers": (
        "<|system|>Use the available tools safely.<|end_of_turn|>"
        "<|user|>Inspect the service and repair it.<|end_of_turn|>"
        "<|assistant|><|think_start|>Check health, logs, then configuration."
        "<|think_end|><|tool_call_start|>{\"name\":\"service_status\","
        "\"arguments\":{\"name\":\"demo\"}}<|tool_call_end|>",
        "<|tool_result_start|>{\"state\":\"running\",\"ready\":true}"
        "<|tool_result_end|><|final_start|>The service is healthy.<|final_end|>"
        "<|end_of_turn|>",
    ),
}


def _count_hf(tokenizer: Tokenizer, text: str) -> int:
    return len(tokenizer.encode(text, add_special_tokens=False).ids)


def _count_tiktoken(encoding: tiktoken.Encoding, text: str) -> int:
    return len(encoding.encode(text, disallowed_special=()))


def benchmark(tokenizer_path: Path) -> dict:
    tokenizer_json = tokenizer_path / "tokenizer.json" if tokenizer_path.is_dir() else tokenizer_path
    tokenizer = Tokenizer.from_file(str(tokenizer_json))
    encodings = {
        name: tiktoken.get_encoding(name)
        for name in ("r50k_base", "cl100k_base", "o200k_base")
    }
    tokenizers = {"tr_hash_agentic_32k": lambda text: _count_hf(tokenizer, text)}
    tokenizers.update(
        {name: (lambda text, enc=encoding: _count_tiktoken(enc, text)) for name, encoding in encodings.items()}
    )

    categories: dict[str, dict] = {}
    totals = {name: defaultdict(int) for name in tokenizers}
    natural = {name: defaultdict(int) for name in tokenizers}
    roundtrip = {}
    for category, samples in SUITE.items():
        texts = tuple(unicodedata.normalize("NFC", sample) for sample in samples)
        characters = sum(len(text) for text in texts)
        utf8_bytes = sum(len(text.encode("utf-8")) for text in texts)
        counts = {name: sum(counter(text) for text in texts) for name, counter in tokenizers.items()}
        categories[category] = {
            "characters": characters,
            "utf8_bytes": utf8_bytes,
            "tokens": counts,
        }
        for name, count in counts.items():
            totals[name]["characters"] += characters
            totals[name]["utf8_bytes"] += utf8_bytes
            totals[name]["tokens"] += count
            if category != "agentic_markers":
                natural[name]["characters"] += characters
                natural[name]["utf8_bytes"] += utf8_bytes
                natural[name]["tokens"] += count
        roundtrip[category] = all(
            tokenizer.decode(
                tokenizer.encode(text, add_special_tokens=False).ids,
                skip_special_tokens=False,
            )
            == text
            for text in texts
        )

    def summarize(values):
        return {
            name: {
                **dict(metric),
                "characters_per_token": metric["characters"] / metric["tokens"],
                "bytes_per_token": metric["utf8_bytes"] / metric["tokens"],
                "vocab_size": (
                    tokenizer.get_vocab_size(with_added_tokens=True)
                    if name == "tr_hash_agentic_32k"
                    else encodings[name].n_vocab
                ),
            }
            for name, metric in values.items()
        }

    return {
        "schema": "tr-hash-agentic-tokenizer-benchmark-v1",
        "suite": "fixed multilingual-agentic suite",
        "categories": categories,
        "natural_text_total": summarize(natural),
        "including_native_markers_total": summarize(totals),
        "tr_hash_roundtrip": roundtrip,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("tokenizer")
    parser.add_argument("--output")
    args = parser.parse_args()
    report = benchmark(Path(args.tokenizer))
    payload = json.dumps(report, indent=2, ensure_ascii=False) + "\n"
    if args.output:
        Path(args.output).write_text(payload, encoding="utf-8")
    print(payload, end="")


if __name__ == "__main__":
    main()