"""
MotherCoreTokenizer — HF-compatible wrapper around MotherTokenizer.

Uses the SentencePiece model at /home/ckenna/model_archive/mother.model
which is the sovereign MediaStream AI tokenizer — NOT a copy of any
external tokenizer.
"""
import os
import sys
from typing import List, Optional


from transformers import PreTrainedTokenizer

import sentencepiece as spm


class MotherCoreTokenizer(PreTrainedTokenizer):
    """
    HF-compatible tokenizer wrapping the MOTHER CORE SentencePiece model.

    Uses the sovereign tokenizer at /home/ckenna/model_archive/mother.model.
    """
    vocab_files_names = {"vocab_file": "tokenizer.model"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file: str = "tokenizer.model",
        bos_token: str = "<s>",
        eos_token: str = "</s>",
        pad_token: str = "<pad>",
        unk_token: str = "<unk>",
        add_bos_token: bool = False,
        add_eos_token: bool = False,
        **kwargs,
    ):
        self.vocab_file = vocab_file
        self.sp_model = spm.SentencePieceProcessor()
        self.sp_model.load(vocab_file)
        self.add_bos_token = add_bos_token
        self.add_eos_token = add_eos_token

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            pad_token=pad_token,
            unk_token=unk_token,
            add_bos_token=add_bos_token,
            add_eos_token=add_eos_token,
            **kwargs,
        )

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

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

    def _tokenize(self, text: str) -> List[str]:
        return self.sp_model.encode(text, out_type=str)

    def _convert_token_to_id(self, token: str) -> int:
        return self.sp_model.piece_to_id(token)

    def _convert_id_to_token(self, index: int) -> str:
        return self.sp_model.id_to_piece(index)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        return self.sp_model.decode(tokens)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        bos = [self.bos_token_id] if self.add_bos_token else []
        eos = [self.eos_token_id] if self.add_eos_token else []
        out = bos + token_ids_0 + eos
        if token_ids_1 is not None:
            out = out + bos + token_ids_1 + eos
        return out

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