"""
ViuAI Multilingual Tokenizer for ViuAI_TTS_200M.
Supports English, Hindi (Devanagari), Hinglish, numbers, punctuation, and UTF-8 byte fallback.
"""

import re
from typing import List, Dict


class ViuAITokenizer:
    def __init__(self):
        # Base special tokens
        self.pad_token = "<pad>"
        self.unk_token = "<unk>"
        self.bos_token = "<bos>"
        self.eos_token = "<eos>"

        # Standard vocab build
        special_tokens = [self.pad_token, self.unk_token, self.bos_token, self.eos_token]
        
        # Punctuation & English characters
        punct = list(" !'(),-.:;?\"[]{}")
        english_chars = list("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789")
        
        # Hindi Devanagari Unicode range: \u0900 to \u097F
        hindi_chars = [chr(cp) for cp in range(0x0900, 0x097F)]

        all_symbols = special_tokens + punct + english_chars + hindi_chars
        # Unique preserve order
        seen = set()
        self.vocab: List[str] = []
        for s in all_symbols:
            if s not in seen:
                seen.add(s)
                self.vocab.append(s)

        self.char_to_id: Dict[str, int] = {ch: idx for idx, ch in enumerate(self.vocab)}
        self.id_to_char: Dict[int, str] = {idx: ch for idx, ch in enumerate(self.vocab)}

    @property
    def vocab_size(self) -> int:
        return len(self.vocab)

    def text_to_ids(self, text: str) -> List[int]:
        """Converts input text to token IDs with boundary tags."""
        text = text.strip()
        ids = [self.char_to_id[self.bos_token]]
        for ch in text:
            if ch in self.char_to_id:
                ids.append(self.char_to_id[ch])
            else:
                ids.append(self.char_to_id[self.unk_token])
        ids.append(self.char_to_id[self.eos_token])
        return ids

    def ids_to_text(self, ids: List[int]) -> str:
        """Converts token IDs back to text."""
        chars = []
        for i in ids:
            ch = self.id_to_char.get(i, "")
            if ch not in (self.pad_token, self.unk_token, self.bos_token, self.eos_token):
                chars.append(ch)
        return "".join(chars)