"""
MnBPETokenizer - HuggingFace PreTrainedTokenizer for the Mongolian BPE model.

Special token layout (matches training in mn_tokenizer_patch.py):
    id 0 -> <pad>
    id 1 -> <s>   BOS
    id 2 -> </s>  EOS / PAD
    id 3+ -> BPE pieces (offset = 3)
"""
import os, shutil
from typing import Dict, List, Optional, Tuple
import sentencepiece as spm
from transformers import PreTrainedTokenizer

VOCAB_FILES_NAMES = {"vocab_file": "mn_bpe.model"}
_SPECIAL = {0: "<pad>", 1: "<s>", 2: "</s>"}
_OFFSET = 3


class MnBPETokenizer(PreTrainedTokenizer):
    vocab_files_names = VOCAB_FILES_NAMES
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file, bos_token="<s>", eos_token="</s>",
                 unk_token="<unk>", pad_token="</s>",
                 sp_model_kwargs=None, **kwargs):
        self.sp_model_kwargs = sp_model_kwargs or {}
        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
        self.sp_model.Load(vocab_file)
        self.vocab_file = vocab_file
        super().__init__(bos_token=bos_token, eos_token=eos_token,
                         unk_token=unk_token, pad_token=pad_token,
                         sp_model_kwargs=sp_model_kwargs, **kwargs)

    @property
    def vocab_size(self):
        return self.sp_model.get_piece_size() + _OFFSET

    def get_vocab(self):
        v = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
        v.update(self.added_tokens_encoder)
        return v

    def _tokenize(self, text):
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token):
        rev = {v: k for k, v in _SPECIAL.items()}
        if token in rev:
            return rev[token]
        return self.sp_model.piece_to_id(token) + _OFFSET

    def _convert_id_to_token(self, index):
        if index in _SPECIAL:
            return _SPECIAL[index]
        return self.sp_model.id_to_piece(index - _OFFSET)

    def convert_tokens_to_string(self, tokens):
        return self.sp_model.decode(tokens)

    def save_vocabulary(self, save_directory, filename_prefix=None):
        if not os.path.isdir(save_directory):
            return ()
        fname = VOCAB_FILES_NAMES["vocab_file"]
        if filename_prefix:
            fname = f"{filename_prefix}-{fname}"
        out = os.path.join(save_directory, fname)
        if os.path.abspath(self.vocab_file) != os.path.abspath(out):
            shutil.copyfile(self.vocab_file, out)
        return (out,)

    def decode_ids(self, ids, skip_special=True):
        """Decode token ids to text, matching training decode logic."""
        pieces = []
        for i in ids:
            i = int(i)
            if i == self.eos_token_id:
                break
            if skip_special and i < _OFFSET:
                continue
            pieces.append(i - _OFFSET)
        return self.sp_model.decode(pieces) if pieces else ""