from tokenizers import Tokenizer as HFTokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
import os

class BPETokenizerWrapper():
    def __init__(self, hf_tokenizer):
        self.tok = hf_tokenizer
        self.word_to_idx = hf_tokenizer.get_vocab()
        self.idx_to_word = {v: k for k, v in self.word_to_idx.items()}

    def __len__(self):
        return self.tok.get_vocab_size()

    def encode(self, text): 
        ids = self.tok.encode(text, add_special_tokens=False).ids
        bos = self.word_to_idx["<bos>"]
        eos = self.word_to_idx["<eos>"]
        return [bos] + ids + [eos]

    def decode(self, ids):
        special_ids = {self.word_to_idx[t] for t in ["<pad>", "<bos>", "<eos>", "<unk>"]}
        filtered = [i for i in ids if i not in special_ids]
        return self.tok.decode(filtered)

def build_bpe_tokenizer(combined_texts, vocab_size=32000, save_path="fineweb_bpe.json"):
    if os.path.exists(save_path):
        tokenizer_bpe = HFTokenizer.from_file(save_path)
    else:
        tokenizer_bpe = HFTokenizer(BPE(unk_token="<unk>"))
        tokenizer_bpe.pre_tokenizer = ByteLevel()
        tokenizer_bpe.decoder = ByteLevelDecoder()

        trainer = BpeTrainer(
            vocab_size=vocab_size,
            special_tokens=["<pad>", "<unk>", "<bos>", "<eos>"],
            initial_alphabet=ByteLevel.alphabet(),
        )

        tokenizer_bpe.train_from_iterator(combined_texts, trainer=trainer)
        tokenizer_bpe.save(save_path)

    return BPETokenizerWrapper(tokenizer_bpe) 

def load_bpe_tokenizer(path):
    if not os.path.exists(path):
        raise FileNotFoundError(f"{path} missing. Run prepare_fineweb.py first.")
    return BPETokenizerWrapper(HFTokenizer.from_file(str(path)))