"""
Tokenizer Builder for 5M Terminal & Multilingual LM
Trains a custom BPE Tokenizer with vocab_size = 4096.
"""
import os
import json
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors
from transformers import PreTrainedTokenizerFast
from terminal_dataset import generate_terminal_dataset

def build_tokenizer(output_dir="tokenizer_5m", vocab_size=4096):
    os.makedirs(output_dir, exist_ok=True)
    print(f"[Tokenizer] Training custom BPE Tokenizer (vocab_size={vocab_size})...")

    # Initialize Byte-Level BPE tokenizer
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)

    special_tokens = ["<pad>", "<s>", "</s>", "<unk>", "<cmd>", "</cmd>", "<user>", "<assistant>"]

    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=special_tokens,
        min_frequency=2,
        show_progress=True
    )

    # Prepare training text samples
    print("[Tokenizer] Generating training corpus samples...")
    terminal_text = generate_terminal_dataset(num_samples=25000)

    # General English text samples (30%)
    english_samples = [
        "The quick brown fox jumps over the lazy dog. Artificial intelligence and machine learning models continue to advance.",
        "System architecture consists of frontend user interfaces, backend APIs, microservices, and high performance databases.",
        "In computer science, algorithms process data structures to solve complex mathematical and computational problems efficiently.",
        "Operating systems manage hardware resources, memory allocation, process scheduling, and file system permissions.",
        "Software development requires version control, continuous integration, testing, code review, and automated deployment.",
        "Deep learning neural networks utilize backpropagation, gradient descent, activation functions, and transformer attention layers.",
        "Linux kernel modules provide driver support for hardware peripherals, network interfaces, and storage volume controllers."
    ] * 2000

    # General Multilingual text samples (70%)
    multilingual_samples = [
        # Indonesian
        "Model bahasa buatan ini dilatih untuk memahami perintah terminal Linux dan teks serbaguna secara efisien.",
        "Sistem operasi Linux menyediakan antarmuka baris perintah yang sangat kuat untuk mengelola berkas dan proses.",
        "Pengembangan perangkat lunak membutuhkan manajemen kode, pengujian otomatis, dan infrastruktur komputasi awan.",
        # Spanish
        "El modelo de lenguaje artificial aprende comandos de terminal Linux y procesamiento de texto en varios idiomas.",
        "Los sistemas de computación moderna utilizan controladores de memoria y procesamiento paralelo en GPU.",
        # French
        "Le modèle linguistique est conçu pour traiter les commandes système et le texte multilingue rapidement.",
        "L'apprentissage automatique et le traitement du langage naturel permettent des interactions intelligentes.",
        # German
        "Das künstliche Sprachmodell lernt Befehle für die Linux-Konsole und mehrsprachige Textverarbeitung.",
        "Moderne Algorithmen optimieren die Datenverarbeitung und die Ausführung von Skripten auf Servern."
    ] * 2000

    training_corpus = [terminal_text] + english_samples + multilingual_samples

    # Train tokenizer
    tokenizer.train_from_iterator(training_corpus, trainer=trainer)

    # Wrap in Transformers PreTrainedTokenizerFast
    fast_tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=tokenizer,
        bos_token="<s>",
        eos_token="</s>",
        pad_token="<pad>",
        unk_token="<unk>",
        mask_token="<mask >",
        additional_special_tokens=["<cmd>", "</cmd>", "<user>", "<assistant>"]
    )

    fast_tokenizer.save_pretrained(output_dir)
    print(f"[Tokenizer] Saved tokenizer to '{output_dir}'. Vocab size: {len(fast_tokenizer)}")

    # Test encoding/decoding
    test_str = "cd .. && ls -la # Check directory\nModel bahasa 5M parameter."
    encoded = fast_tokenizer.encode(test_str)
    decoded = fast_tokenizer.decode(encoded)
    print(f"\n[Test Encoding]\nInput:  {test_str}\nTokens: {encoded[:15]}...\nDecoded: {decoded}")

    return fast_tokenizer

if __name__ == "__main__":
    build_tokenizer()