import json
from pathlib import Path

from tokenizers import Tokenizer, models, trainers, pre_tokenizers, decoders

VOCAB_SIZE = 32000
SPECIAL_TOKENS = ["<pad>", "<unk>", "<bos>", "<eos>"]


class BPETokenizer:
    def __init__(self, tokenizer: Tokenizer):
        self._tok = tokenizer

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

    def encode(self, text: str) -> list[int]:
        return self._tok.encode(text).ids

    def decode(self, ids) -> str:
        return self._tok.decode(list(ids))

    def save(self, path: Path):
        self._tok.save(str(path))

    @classmethod
    def load(cls, path: Path) -> "BPETokenizer":
        return cls(Tokenizer.from_file(str(path)))

    @classmethod
    def build_from_text(cls, text: str, vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
        tok = Tokenizer(models.BPE(unk_token="<unk>"))
        tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
        tok.decoder = decoders.ByteLevel()
        trainer = trainers.BpeTrainer(
            vocab_size=vocab_size,
            special_tokens=SPECIAL_TOKENS,
            min_frequency=2,
        )
        tok.train_from_iterator(_chunk(text), trainer=trainer)
        return cls(tok)

    @classmethod
    def build_from_files(cls, paths: list[Path], vocab_size: int = VOCAB_SIZE) -> "BPETokenizer":
        tok = Tokenizer(models.BPE(unk_token="<unk>"))
        tok.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)
        tok.decoder = decoders.ByteLevel()
        trainer = trainers.BpeTrainer(
            vocab_size=vocab_size,
            special_tokens=SPECIAL_TOKENS,
            min_frequency=2,
        )
        tok.train([str(p) for p in paths], trainer=trainer)
        return cls(tok)


def _chunk(text: str, size: int = 1_000_000):
    for i in range(0, len(text), size):
        yield text[i:i + size]


CharTokenizer = BPETokenizer