"""Sansar Sanskrit tokenizer for transformers: Devanagari in, token ids out, Devanagari back.

The SentencePiece model (tokenizer.model, the same file as MuseMesh/sansar-sanskrit-tokenizer v0.1.0) works on SLP1, a
lossless ASCII spelling of Devanagari. This class converts with the bundled translit.py, so callers only see Devanagari.
Latin-script runs (English words, IAST) are fenced with two private-use marks so they come back unchanged
(``fence_latin=True``, the default). Training and evaluation text was NOT fenced (Latin letters went through the SLP1
step as they were), so for text with Latin runs ``fence_latin=False`` reproduces the training token stream exactly,
at the price of Latin coming back as Devanagari on decode. Pure Devanagari text gives the same ids either way.

No BOS/EOS is added to encoded text: the models were trained on records separated by </s> (id 2), with no BOS.

Code licence: Apache-2.0 (LICENSE-CODE). tokenizer.model: CC BY-NC 4.0 (LICENSE).
"""
from __future__ import annotations

import os
import re
import shutil

import sentencepiece as spm
from transformers import PreTrainedTokenizer

try:
    from .translit import NFC, deva_to_slp1, slp1_to_deva
except ImportError:  # imported as a plain module next to translit.py
    from translit import NFC, deva_to_slp1, slp1_to_deva

VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}

_OPEN, _CLOSE = "", ""   # fence for Latin runs; byte-fallback tokens in the model
_LATIN = re.compile(r"(?:[A-Za-zÀ-ɏḀ-ỿ][̀-ͯ]*)+")
_FENCED = re.compile(re.escape(_OPEN) + r"(.*?)" + re.escape(_CLOSE), re.S)


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

    def __init__(self, vocab_file, unk_token="<unk>", bos_token="<s>", eos_token="</s>", pad_token="</s>",
                 clean_up_tokenization_spaces=False, fence_latin=True, **kwargs):
        self.vocab_file = vocab_file
        self.fence_latin = fence_latin
        self.sp = spm.SentencePieceProcessor(model_file=vocab_file)
        kwargs.setdefault("padding_side", "left")    # left padding for batched generation
        super().__init__(unk_token=unk_token, bos_token=bos_token, eos_token=eos_token, pad_token=pad_token,
                         clean_up_tokenization_spaces=clean_up_tokenization_spaces, fence_latin=fence_latin, **kwargs)

    # ------------------------------------------------------------------ transliteration
    def to_slp1(self, text: str) -> str:
        """Devanagari -> the model's input: SLP1, with Latin runs fenced and kept as they are (if fence_latin)."""
        text, out, pos = NFC(text), [], 0
        if not self.fence_latin:
            return deva_to_slp1(text)[0]
        for m in _LATIN.finditer(text):
            out.append(deva_to_slp1(text[pos:m.start()])[0])
            out.append(_OPEN + m.group() + _CLOSE)
            pos = m.end()
        out.append(deva_to_slp1(text[pos:])[0])
        return "".join(out)

    @staticmethod
    def from_slp1(s: str) -> str:
        out, pos = [], 0
        for m in _FENCED.finditer(s):
            out.append(slp1_to_deva(s[pos:m.start()])[0])
            out.append(m.group(1))
            pos = m.end()
        out.append(slp1_to_deva(s[pos:])[0])
        return "".join(out)

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

    def get_vocab(self):
        vocab = {self.sp.id_to_piece(i): i for i in range(self.sp.get_piece_size())}
        vocab.update(self.added_tokens_encoder)
        return vocab

    def _tokenize(self, text, **kwargs):
        """Text -> SentencePiece pieces in SLP1 (the pieces the model sees)."""
        return self.sp.encode(self.to_slp1(text), out_type=str)

    def _convert_token_to_id(self, token):
        return self.sp.piece_to_id(token)

    def _convert_id_to_token(self, index):
        return self.sp.id_to_piece(int(index))

    def convert_tokens_to_string(self, tokens):
        """Pieces -> Devanagari text (Latin runs restored). Special tokens are kept as written."""
        specials = set(self.all_special_tokens)
        out, buf = [], []
        for t in tokens:
            if t in specials:
                if buf:
                    out.append(self.from_slp1(self.sp.decode_pieces(buf)))
                    buf = []
                out.append(t)
            else:
                buf.append(t)
        if buf:
            out.append(self.from_slp1(self.sp.decode_pieces(buf)))
        return "".join(out)

    # ------------------------------------------------------------------ special tokens: none added
    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        return list(token_ids_0) + (list(token_ids_1) if token_ids_1 is not None 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 super().get_special_tokens_mask(token_ids_0, token_ids_1, already_has_special_tokens=True)
        return [0] * (len(token_ids_0) + (len(token_ids_1) if token_ids_1 is not None else 0))

    def num_special_tokens_to_add(self, pair: bool = False) -> int:
        return 0

    def save_vocabulary(self, save_directory, filename_prefix=None):
        if not os.path.isdir(save_directory):
            os.makedirs(save_directory, exist_ok=True)
        out = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "tokenizer.model")
        if os.path.abspath(self.vocab_file) != os.path.abspath(out):
            shutil.copyfile(self.vocab_file, out)
        return (out,)