#!/usr/bin/env python3
# ============================================================
# ZabaanAI-v2: Tokenizer Improvement Script
# Adds Pakistan language tokens to XLM-RoBERTa tokenizer
# Run BEFORE training for better Urdu/Sindhi/Pashto tokenization
# ============================================================
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel, WhitespaceSplit, DigitsPreTokenizer
from tokenizers.processors import TemplateProcessing
from tokenizers.normalizers import NFD, StripAccents, Sequence as NormSeq
from tokenizers.decoders import BPEDecoder
import os, glob, json

OUT_DIR  = '/app/data/tokenizer_pakistan'
OUT_NAME = 'zabaanai_tokenizer'
os.makedirs(OUT_DIR, exist_ok=True)

# ── Custom Normalizer for Pakistan scripts ────────────────────
class PakistanNormalizer:
    def __init__(self):
        # Urdu Nukta handling, Urdu-Punjabi diacritic normalization
        self.replacements = [
            ('\u06C1\u06C2', '\u06C1'),  # ہے → ہ (heh + shadda → heh)
            ('\u06A9', '\u06A9'),          # Persian keh = Arabic keh
            ('\u06AF', '\u06AF'),          # Persian geh = Arabic geh
            # Normalize Persian/Urdu variants to Arabic script
            ('\u067E', '\u067E'),          # Persian peh
            ('\u06CC', '\u06CC'),          # Persian yeh (ی)
            ('\u06D2', '\u06D2'),          # Persian yeh barree
            # Standardize diacritics
            ('\u064B', ''),   # Fathatan
            ('\u064C', ''),   # Dammatan  
            ('\u064D', ''),   # Kasratan
            ('\u0650', ''),   # Kasra
            ('\u064E', ''),   # Fatha
            ('\u064F', ''),   # Damma
            ('\u0652', ''),   # Sukun
        ]

    def normalize(self, text: str) -> str:
        for old, new in self.replacements:
            text = text.replace(old, new)
        return text

# ── Build BPE tokenizer with Pakistan language data ───────────
def train_pakistan_tokenizer():
    print('Training Pakistan BPE tokenizer...')

    # Start from raw BPE model
    tokenizer = Tokenizer(BPE(unk_token='<unk>'))

    # Normalizer: NFD + custom + strip accents (for Arabic)
    tokenizer.normalizer = NormSeq([NFD(), PakistanNormalizer()])

    # Pre-tokenizer: bytes first, then whitespace, then digits
    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=False)

    # Decoder
    tokenizer.decoder = BPEDecoder(suffix='</w>')

    # Trainer with Pakistan language settings
    trainer = BpeTrainer(
        vocab_size=250_000,        # Large vocab for multilingual
        min_frequency=2,
        max_token_length=100,
        show_progress=True,
        special_tokens=[
            '<pad>', '<unk>', '<s>', '</s>', '<eot>', '<eop>',
            '<|endoftext|>', '<|im_start|>', '<|im_end|>',
        ],
        initial_alphabet=ByteLevel.alphabet(),
    )

    # Collect all text files
    text_files = glob.glob('/app/data/raw/**/*.txt', recursive=True)
    text_files += glob.glob('/app/data/raw/**/*.jsonl', recursive=True)

    if not text_files:
        print('No text files found in /app/data/raw/. Training on empty corpus.')
        print('  → Will download sample data or use pretrained tokenizer.')

    print(f'  Found {len(text_files)} text files for tokenizer training')

    # Train tokenizer
    if text_files:
        tokenizer.train_from_iterator(
            line_reader(text_files),
            trainer=trainer,
            length=sum(1 for _ in file_reader(text_files)),
        )
    else:
        print('  ⚠ No training files — using default pretrained tokenizer')
        return None

    # Post-processing: ChatML template
    tokenizer.post_processor = TemplateProcessing(
        single='$A <|im_end|>',
        pair='$A <|im_end|> $B:0 <|im_end|>',
        special_tokens=[
            ('<|im_start|>', 1),
            ('<|im_end|>', 2),
        ],
    )

    return tokenizer

def line_reader(files):
    for f in files:
        with open(f, 'r', encoding='utf-8', errors='ignore') as fp:
            for line in fp:
                line = line.strip()
                if line:
                    yield line

def file_reader(files):
    for f in files:
        with open(f, 'r', encoding='utf-8', errors='ignore') as fp:
            for line in fp:
                yield line

# ── Main ──────────────────────────────────────────────────────
tokenizer = train_pakistan_tokenizer()
if tokenizer:
    tokenizer.save(os.path.join(OUT_DIR, f'{OUT_NAME}.json'))
    tokenizer.model.save(OUT_DIR, OUT_NAME)
    print(f'Saved tokenizer to {OUT_DIR}')
    print(f'  Vocab size: ~250K (Pakistan-optimized)')
else:
    print('Using default tokenizer (XLM-RoBERTa pretrained)')