"""Thin wrapper around a SentencePiece model with fixed special-token ids."""
from pathlib import Path


class SPTokenizer:
    PAD_ID, UNK_ID, BOS_ID, EOS_ID = 0, 1, 2, 3

    def __init__(self, model_path):
        import sentencepiece as spm

        model_path = Path(model_path)
        if not model_path.exists():
            raise FileNotFoundError(f"Tokenizer model not found: {model_path}")
        self.sp = spm.SentencePieceProcessor()
        self.sp.load(str(model_path))
        got = (self.sp.pad_id(), self.sp.unk_id(), self.sp.bos_id(), self.sp.eos_id())
        want = (self.PAD_ID, self.UNK_ID, self.BOS_ID, self.EOS_ID)
        if got != want:
            raise ValueError(f"Tokenizer special ids {got} != expected {want}. "
                             "Re-train the tokenizer with tokenizer/train_tokenizer.py.")

    @property
    def vocab_size(self) -> int:
        return self.sp.get_piece_size()

    def encode(self, text, add_bos=False, add_eos=False):
        ids = self.sp.encode(text, out_type=int)
        if add_bos:
            ids = [self.BOS_ID] + ids
        if add_eos:
            ids = ids + [self.EOS_ID]
        return ids

    def pieces(self, text):
        return self.sp.encode(text, out_type=str)

    def decode(self, ids):
        ids = [int(i) for i in ids if int(i) not in (self.PAD_ID, self.BOS_ID, self.EOS_ID)]
        return self.sp.decode(ids)