#!/usr/bin/env python3
"""Generate C++ header with Chinese character → pinyin token mapping."""
import sys, os
REPO_DIR = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, REPO_DIR)

from pypinyin.pinyin_dict import pinyin_dict
from pypinyin.contrib.tone_convert import to_initials, to_finals_tone3

TOKEN_FILE = os.path.join(REPO_DIR, 'resources', 'zipvoice_hf', 'zipvoice', 'tokens.txt')
OUTPUT = os.path.join(REPO_DIR, 'cpp', 'src', 'pinyin_table.hpp')

# Load token2id
token2id = {}
with open(TOKEN_FILE) as f:
    for line in f:
        if '\t' in line:
            t, tid = line.rstrip().split('\t')
            token2id[t] = int(tid)

# Build char → token IDs
char_entries = []
for codepoint, pinyins in pinyin_dict.items():
    ch = chr(codepoint)
    py = pinyins.split(',')[0]
    initial = to_initials(py, strict=False)
    final = to_finals_tone3(py, strict=False, neutral_tone_with_five=True)
    if not final: continue

    ini_token = initial + '0' if initial else ''
    tokens = []
    if ini_token and ini_token in token2id:
        tokens.append(token2id[ini_token])
    if final in token2id:
        tokens.append(token2id[final])

    if tokens:
        char_entries.append((codepoint, tokens))

# Sort by codepoint for binary search
char_entries.sort(key=lambda x: x[0])

# Generate C++ header
with open(OUTPUT, 'w') as f:
    f.write('''// Auto-generated Chinese character → token ID mapping
// Source: pypinyin.pinyin_dict, tokens.txt
// Generated by gen_pinyin_table.py

#pragma once
#include <cstdint>
#include <string>
#include <vector>

struct PinyinEntry {
    uint32_t codepoint;
    int16_t  token1;
    int16_t  token2;  // -1 if only one token
};

// Sorted by codepoint, use binary search
static const PinyinEntry PINYIN_TABLE[] = {
''')

    for codepoint, tokens in char_entries:
        t1 = tokens[0]
        t2 = tokens[1] if len(tokens) > 1 else -1
        f.write(f'    {{0x{codepoint:04X}, {t1}, {t2}}},  // {chr(codepoint)}\n')

    f.write(f'''}};

static const int PINYIN_TABLE_SIZE = {len(char_entries)};

// Lookup tokens for a Chinese character. Returns empty vector if not found.
inline std::vector<int> pinyin_lookup(uint32_t codepoint) {{
    // Binary search
    int lo = 0, hi = PINYIN_TABLE_SIZE - 1;
    while (lo <= hi) {{
        int mid = (lo + hi) / 2;
        if (PINYIN_TABLE[mid].codepoint == codepoint) {{
            std::vector<int> result;
            result.push_back(PINYIN_TABLE[mid].token1);
            if (PINYIN_TABLE[mid].token2 >= 0)
                result.push_back(PINYIN_TABLE[mid].token2);
            return result;
        }}
        if (PINYIN_TABLE[mid].codepoint < codepoint)
            lo = mid + 1;
        else
            hi = mid - 1;
    }}
    return {{}};
}}
''')

print(f"Generated {OUTPUT}: {len(char_entries)} entries")