# -*- coding: utf-8 -*-
from pathlib import Path
from typing import Iterable, List, Union

import sentencepiece as spm


class SentencepiecesTokenizer:
    def __init__(self, bpemodel: Union[Path, str], **kwargs):
        self.bpemodel = str(bpemodel)
        self.sp = None
        self._build()

    def __repr__(self):
        return f'{self.__class__.__name__}(model="{self.bpemodel}")'

    def _build(self):
        if self.sp is None:
            self.sp = spm.SentencePieceProcessor()
            self.sp.load(self.bpemodel)

    def text2tokens(self, line: str) -> List[str]:
        self._build()
        return self.sp.EncodeAsPieces(line)

    def tokens2text(self, tokens: Iterable[str]) -> str:
        self._build()
        return self.sp.DecodePieces(list(tokens))

    def encode(self, line: str) -> List[int]:
        self._build()
        return self.sp.EncodeAsIds(line)

    def decode(self, line: List[int]):
        self._build()
        return self.sp.DecodeIds(line)

    def get_vocab_size(self):
        self._build()
        return self.sp.GetPieceSize()

    def ids2tokens(self, *args, **kwargs):
        return self.decode(*args, **kwargs)

    def tokens2ids(self, *args, **kwargs):
        return self.encode(*args, **kwargs)