#!/usr/bin/env python3
"""Emit `src/g2p/vocab.rs` from Kokoro's `config.json`.

    python3 gen_vocab.py config.json ../src/g2p/vocab.rs
"""

import json
import sys


def main():
    src, dst = sys.argv[1], sys.argv[2]
    vocab = json.load(open(src))["vocab"]
    pairs = sorted(vocab.items(), key=lambda kv: kv[0])

    lines = [
        "// Generated by tools/gen_vocab.py from Kokoro's config.json. Do not edit.",
        "",
        "/// Phoneme symbol -> token id, sorted by symbol for binary search.",
        f"pub const VOCAB: [(char, u16); {len(pairs)}] = [",
    ]
    for symbol, token in pairs:
        # Always escape: the vocab contains `'`, `"`, a combining tilde and IPA.
        escaped = f"\\u{{{ord(symbol):04x}}}"
        shown = symbol if symbol.strip() and symbol != "'" else repr(symbol)
        lines.append(f"    ('{escaped}', {token}), // {shown}")
    lines.append("];")
    lines.append("")
    lines.append("/// Look up a phoneme symbol's token id.")
    lines.append("pub fn token(symbol: char) -> Option<u16> {")
    lines.append("    VOCAB")
    lines.append("        .binary_search_by_key(&symbol, |(sym, _)| *sym)")
    lines.append("        .ok()")
    lines.append("        .map(|index| VOCAB[index].1)")
    lines.append("}")
    lines.append("")

    with open(dst, "w") as out:
        out.write("\n".join(lines))
    print(f"{dst}: {len(pairs)} symbols")


if __name__ == "__main__":
    main()