# /// script
# requires-python = ">=3.11,<3.13"
# dependencies = ["transformers==4.57.6"]
# ///
"""
Emit expected (text, token_ids) pairs for the Qwen3-ASR byte-level BPE
encoder parity test.

Reads the tokenizer from a Qwen3-ASR checkpoint directory (argv[1]) and
prints a C++ snippet the test binary #includes. Two sections:

  k_bpe_fixtures[]          strings with no Qwen3-ASR special tokens;
                            encodes through our Tokenizer::encode() and
                            must match HF's encode(add_special_tokens=
                            False) token-for-token.

  k_lang_prefix_fixtures[]  the 30 "language {Name}" prefixes (without
                            the <asr_text> tag) plus the expected
                            full-prefix ids (including <asr_text>).
                            The test calls encode_language_prefix()
                            from the qwen3_asr internals and compares
                            against the full-prefix ids.

Usage:
  uv run scripts/tokenizer-parity-fixture.py \\
      ~/sandboxes/transcribe/models/Qwen3-ASR-0.6B \\
      > tests/fixtures/qwen3_asr_bpe_parity.inc
"""
import sys
from pathlib import Path

from transformers import AutoTokenizer

# Publisher canonical names, in the exact order of qwen_asr.inference.
# utils.SUPPORTED_LANGUAGES.
SUPPORTED_LANGUAGES = [
    "Chinese", "English", "Cantonese", "Arabic", "German", "French",
    "Spanish", "Portuguese", "Indonesian", "Italian", "Korean",
    "Russian", "Thai", "Vietnamese", "Japanese", "Turkish", "Hindi",
    "Malay", "Dutch", "Swedish", "Danish", "Finnish", "Polish",
    "Czech", "Filipino", "Persian", "Greek", "Romanian", "Hungarian",
    "Macedonian",
]

# BCP-47 codes aligned with SUPPORTED_LANGUAGES. Duplicated from the
# converter's LANGUAGE_TO_BCP47; this script is a dev tool so a local
# copy is fine.
BCP47 = [
    "zh", "en", "yue", "ar", "de", "fr", "es", "pt", "id", "it",
    "ko", "ru", "th", "vi", "ja", "tr", "hi", "ms", "nl", "sv",
    "da", "fi", "pl", "cs", "fil", "fa", "el", "ro", "hu", "mk",
]

# Strings WITHOUT <asr_text>. The encoder must reproduce HF exactly.
PURE_BPE = [
    "hello world",
    "it's a test, isn't it?",
    "   trailing   spaces",
    "line1\nline2",
    "digits: 12345 and 678",
    "mixed \u0410\u0411\u0412 and \u65e5\u672c\u8a9e",
    "emoji: \U0001F600 \U0001F389",
] + [f"language {name}" for name in SUPPORTED_LANGUAGES]


def c_escape(s: str) -> str:
    out = []
    for ch in s.encode("utf-8"):
        if ch == 0x22:
            out.append("\\\"")
        elif ch == 0x5C:
            out.append("\\\\")
        elif ch == 0x0A:
            out.append("\\n")
        elif ch == 0x0D:
            out.append("\\r")
        elif ch == 0x09:
            out.append("\\t")
        elif 0x20 <= ch < 0x7F:
            out.append(chr(ch))
        else:
            # Emit `"\xAB" "<next-char>"` to avoid greedy hex-escape
            # merging across adjacent bytes (e.g. "\xC4" "\x8A" not
            # "\xC48A").
            out.append(f"\\x{ch:02x}\" \"")
    return "".join(out)


def emit_array(name: str, items) -> None:
    print(f"constexpr Fixture {name}[] = {{")
    for text, ids in items:
        ids_csv = ", ".join(str(i) for i in ids)
        print(f'    {{ "{c_escape(text)}", {{ {ids_csv} }}, {len(ids)} }},')
    print("};")
    print(f"constexpr size_t {name}_n = sizeof({name}) / sizeof({name}[0]);")
    print()


def main() -> int:
    if len(sys.argv) != 2:
        print(__doc__, file=sys.stderr)
        return 2
    ckpt = Path(sys.argv[1])
    tok = AutoTokenizer.from_pretrained(str(ckpt))

    # Pure-BPE fixtures: text, ids.
    pure = []
    for s in PURE_BPE:
        ids = tok.encode(s, add_special_tokens=False)
        pure.append((s, ids))

    # Language-prefix fixtures: bcp47, pub_name, full_expected_ids.
    lang = []
    for code, name in zip(BCP47, SUPPORTED_LANGUAGES):
        full = f"language {name}<asr_text>"
        ids = tok.encode(full, add_special_tokens=False)
        lang.append((code, name, ids))

    print("// Auto-generated by scripts/tokenizer-parity-fixture.py.")
    print("// Regenerate with:")
    print("//   uv run scripts/tokenizer-parity-fixture.py"
          " <checkpoint-dir>")
    print("//   > tests/fixtures/qwen3_asr_bpe_parity.inc")
    print("//")
    print("// The checkpoint directory is local to each developer's machine")
    print("// so we don't bake the absolute path into the committed header.")
    print()

    emit_array("k_bpe_fixtures", pure)

    print("struct LangFixture {")
    print("    const char * bcp47;")
    print("    const char * pub_name;")
    print("    const int32_t ids[16];")
    print("    size_t n_ids;")
    print("};")
    print("constexpr LangFixture k_lang_prefix_fixtures[] = {")
    for code, name, ids in lang:
        ids_csv = ", ".join(str(i) for i in ids)
        print(f'    {{ "{code}", "{name}", {{ {ids_csv} }}, {len(ids)} }},')
    print("};")
    print("constexpr size_t k_lang_prefix_fixtures_n = "
          "sizeof(k_lang_prefix_fixtures) / sizeof(k_lang_prefix_fixtures[0]);")

    return 0


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