"""
Custom Byte-Pair Encoding (BPE) Tokenizer for Nova 1.0
Built from scratch using modern tokenization principles.
"""

import os
from typing import List, Union
from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders, processors


class Nova1Tokenizer:
    SPECIAL_TOKENS = ["<pad>", "<unk>", "<bos>", "<eos>", "<mask() >"]

    def __init__(self, tokenizer_path: str = None):
        if tokenizer_path and os.path.exists(tokenizer_path):
            self.tokenizer = Tokenizer.from_file(tokenizer_path)
        else:
            self.tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
            self.tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
            self.tokenizer.decoder = decoders.ByteLevel()
            self.tokenizer.post_processor = processors.ByteLevel(trim_offsets=False)

    def train_from_iterator(self, iterator, vocab_size: int = 8192):
        """Train BPE vocabulary on an iterator of text strings."""
        trainer = trainers.BpeTrainer(
            vocab_size=vocab_size,
            special_tokens=self.SPECIAL_TOKENS,
            initial_alphabet=pre_tokenizers.ByteLevel.alphabet()
        )
        self.tokenizer.train_from_iterator(iterator, trainer=trainer)

    def encode(self, text: str, add_bos: bool = True, add_eos: bool = True) -> List[int]:
        """Encode string into list of token IDs."""
        encoding = self.tokenizer.encode(text)
        ids = encoding.ids
        if add_bos:
            ids = [self.bos_id] + ids
        if add_eos:
            ids = ids + [self.eos_id]
        return ids

    def decode(self, ids: List[int], skip_special_tokens: bool = True) -> str:
        """Decode list of token IDs back into string."""
        return self.tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)

    def save(self, path: str):
        """Save tokenizer configuration to JSON file."""
        os.makedirs(os.path.dirname(path), exist_ok=True)
        self.tokenizer.save(path)

    @classmethod
    def load(cls, path: str) -> "Nova1Tokenizer":
        return cls(tokenizer_path=path)

    @property
    def vocab_size(self) -> int:
        return self.tokenizer.get_vocab_size()

    @property
    def pad_id(self) -> int:
        return self.tokenizer.token_to_id("<pad>") or 0

    @property
    def unk_id(self) -> int:
        return self.tokenizer.token_to_id("<unk>") or 1

    @property
    def bos_id(self) -> int:
        return self.tokenizer.token_to_id("<bos>") or 2

    @property
    def eos_id(self) -> int:
        return self.tokenizer.token_to_id("<eos>") or 3

    @property
    def mask_id(self) -> int:
        return self.tokenizer.token_to_id("<mask() >") or 4