#!/usr/bin/env python3
# Copyright 2026 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
"""Generate expected token IDs for tokenizer streaming tests.

This script loads tokenizer.json files and generates expected token ID arrays
that can be embedded directly in C++ tests. The expected values are ground truth
from HuggingFace's tokenizers library.

Usage:
    # Generate expected IDs for all test inputs (requires tokenizers library)
    uv run --with tokenizers python generate_testdata_expected_ids.py

    # Or with pip:
    pip install tokenizers
    python generate_testdata_expected_ids.py
"""

import json
import os
from pathlib import Path

# Import tokenizers library
try:
    from tokenizers import Tokenizer
except ImportError:
    print(
        "Error: tokenizers library not installed.\n"
        "Install with: pip install tokenizers\n"
        "Or run with: uv run --with tokenizers python generate_testdata_expected_ids.py"
    )
    exit(1)


def load_tokenizer_from_json(json_path: Path) -> Tokenizer:
    """Load a tokenizer from a tokenizer.json file."""
    return Tokenizer.from_file(str(json_path))


def encode_and_format(tokenizer: Tokenizer, text: str, name: str) -> str:
    """Encode text and return C++ formatted expected IDs."""
    encoding = tokenizer.encode(text, add_special_tokens=False)
    ids = encoding.ids

    # Format as C++ array
    ids_str = ", ".join(str(id) for id in ids)

    # Escape the text for C++ string literal
    escaped_text = text.replace("\\", "\\\\").replace('"', '\\"').replace("\n", "\\n")

    return f"""  // {name}: "{escaped_text}"
  // Tokens: {encoding.tokens}
  static constexpr iree_tokenizer_token_id_t k{name.replace(" ", "")}Expected[] = {{{ids_str}}};
"""


def main():
    # Find testdata directory
    script_dir = Path(__file__).parent
    testdata_dir = script_dir.parent / "testdata"

    # Test inputs for the bpe_bytelevel_minimal.json tokenizer
    # This tokenizer has merges for "hello" and "world" specifically
    test_cases = [
        ("Hello", "hello"),  # Lowercase gets different treatment
        ("HelloWorld", "hello world"),
        ("SimpleASCII", "The quick brown fox"),
        ("WithSpecial", "hello<|endoftext|>world"),
        ("Unicode", "café résumé"),
        ("CJK", "你好世界"),
        ("Mixed", "Hello 世界!"),
        ("Whitespace", "  hello   world  "),
        ("Newlines", "hello\nworld\n"),
        ("Empty", ""),
        ("SingleChar", "a"),
        ("Numbers", "12345"),
        ("Punctuation", "!@#$%^&*()"),
        ("LongRepeat", "a" * 100),
    ]

    # Load the minimal BPE tokenizer
    tokenizer_path = testdata_dir / "bpe_bytelevel_minimal.json"
    if not tokenizer_path.exists():
        print(f"Error: Tokenizer not found: {tokenizer_path}")
        return

    print(f"Loading tokenizer from: {tokenizer_path}")
    tokenizer = load_tokenizer_from_json(tokenizer_path)

    print("\n// Expected token IDs for streaming encode tests")
    print("// Generated by generate_testdata_expected_ids.py")
    print("// Ground truth from HuggingFace tokenizers library")
    print()

    for name, text in test_cases:
        print(encode_and_format(tokenizer, text, name))

    # Also print a summary table
    print("\n/* Summary table:")
    print(f"{'Name':<20} {'Input':<30} {'Token Count'}")
    print("-" * 60)
    for name, text in test_cases:
        encoding = tokenizer.encode(text, add_special_tokens=False)
        display_text = repr(text)[:28]
        print(f"{name:<20} {display_text:<30} {len(encoding.ids)}")
    print("*/")


if __name__ == "__main__":
    main()
