"""HF tokenizer adapter preserving this project's exact SentencePiece ID mapping."""
from pathlib import Path

import sentencepiece as spm
from transformers import PreTrainedTokenizer


class ClassicalTokenizer(PreTrainedTokenizer):
    vocab_files_names={'vocab_file':'tokenizer.model'}
    model_input_names=['input_ids','attention_mask']

    def __init__(self,vocab_file,add_bos_token=True,add_eos_token=False,**kwargs):
        self.vocab_file=str(vocab_file)
        self.sp_model=spm.SentencePieceProcessor(model_file=self.vocab_file)
        self.marker_ids=[self.sp_model.piece_to_id(f'<0x{b:02X}>') for b in '▁'.encode()]
        self.add_bos_token=add_bos_token;self.add_eos_token=add_eos_token
        for name,i in [('unk_token',0),('bos_token',1),('eos_token',2),('pad_token',3)]:
            kwargs.setdefault(name,self.sp_model.id_to_piece(i))
        kwargs.setdefault('clean_up_tokenization_spaces',False)
        kwargs.setdefault('model_max_length',8192)
        super().__init__(add_bos_token=add_bos_token,add_eos_token=add_eos_token,**kwargs)
        if self.vocab_size!=32000 or (self.unk_token_id,self.bos_token_id,self.eos_token_id,self.pad_token_id)!=(0,1,2,3):
            raise ValueError('Tokenizer IDs were remapped')

    @property
    def vocab_size(self):return self.sp_model.vocab_size()

    def get_vocab(self):return {self.sp_model.id_to_piece(i):i for i in range(self.vocab_size)}

    def raw_encode(self,text):
        ids=[]
        for i,part in enumerate(text.split('▁')):
            if i:ids.extend(self.marker_ids)
            ids.extend(self.sp_model.encode(part))
        return ids

    def tokenize(self,text,**kwargs):
        # Do not reinterpret literal '<s>' etc as injected special tokens.
        return [self.sp_model.id_to_piece(i) for i in self.raw_encode(text)]

    def _tokenize(self,text,**kwargs):return self.tokenize(text,**kwargs)
    def _convert_token_to_id(self,token):return self.sp_model.piece_to_id(token)
    def _convert_id_to_token(self,index):return self.sp_model.id_to_piece(index)
    def convert_tokens_to_string(self,tokens):return self.sp_model.decode(tokens)

    def build_inputs_with_special_tokens(self,token_ids_0,token_ids_1=None):
        if token_ids_1 is not None:raise ValueError('Paired inputs are not supported')
        return ([self.bos_token_id] if self.add_bos_token else [])+list(token_ids_0)+([self.eos_token_id] if self.add_eos_token else [])

    def get_special_tokens_mask(self,token_ids_0,token_ids_1=None,already_has_special_tokens=False):
        if already_has_special_tokens:return [int(i in (0,1,2,3)) for i in token_ids_0]
        return ([1] if self.add_bos_token else [])+[0]*len(token_ids_0)+([1] if self.add_eos_token else [])

    def decode(self,token_ids,skip_special_tokens=False,**kwargs):
        if hasattr(token_ids,'tolist'):token_ids=token_ids.tolist()
        if isinstance(token_ids,int):token_ids=[token_ids]
        if skip_special_tokens:token_ids=[i for i in token_ids if i not in (0,1,2,3)]
        return self.sp_model.decode(token_ids)

    def save_vocabulary(self,save_directory,filename_prefix=None):
        target=Path(save_directory)/((filename_prefix+'-' if filename_prefix else '')+'tokenizer.model')
        target.write_bytes(self.sp_model.serialized_model_proto())
        return (str(target),)


ClassicalTokenizer.register_for_auto_class('AutoTokenizer')