import os
import json
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors

def train_domain_tokenizer(input_text_path: str, output_dir: str, vocab_size: int = 32000):
    os.makedirs(output_dir, exist_ok=True)
    
    # Initialize Byte-Pair Encoding (BPE) Tokenizer
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    tokenizer.decoder = decoders.ByteLevel()
    
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        special_tokens=["<unk>", "<s>", "</s>", "<pad>", "<mask>"],
        min_frequency=2
    )
    
    print(f"Training BPE Tokenizer on {input_text_path} (vocab_size={vocab_size})...")
    tokenizer.train(files=[input_text_path], trainer=trainer)
    
    tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)
    
    tokenizer_path = os.path.join(output_dir, "tokenizer.json")
    tokenizer.save(tokenizer_path)
    print(f"Saved custom domain tokenizer to {tokenizer_path}")

if __name__ == "__main__":
    sample_data_path = "/tmp/sample_domain_corpus.txt"
    if not os.path.exists(sample_data_path):
        with open(sample_data_path, "w") as f:
            f.write("SELECT * FROM users WHERE status = 'active';\n" * 100)
    
    train_domain_tokenizer(sample_data_path, "/home/adminuser/foundational_model/tokenizer", vocab_size=1000)