"""SentencePiece tokenizer wrapper for TaoNet."""

import json
import os
import re
import shutil

from transformers import PreTrainedTokenizer


class TaoNetTokenizer(PreTrainedTokenizer):
    """Transformers-compatible SentencePiece tokenizer for TaoNet."""

    vocab_files_names = {"vocab_file": "tokenizer.model", "special_tokens_file": "tokenizer.special_tokens.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file,
        special_tokens_file=None,
        bos_token="<BOS>",
        eos_token="<EOS>",
        unk_token="<UNK>",
        pad_token="<PAD>",
        additional_special_tokens=None,
        **kwargs,
    ):
        try:
            import sentencepiece as spm
        except ImportError as exc:
            raise ImportError("TaoNetTokenizer requires sentencepiece to be installed") from exc

        # Newer Transformers versions may round-trip this field from tokenizer_config.json
        # as either a dict or a list. TaoNet only needs the token strings here.
        extra_special_tokens = kwargs.pop("extra_special_tokens", None)
        self.vocab_file = vocab_file
        self.special_tokens_file = special_tokens_file
        self.sp_model = spm.SentencePieceProcessor()
        self.sp_model.Load(vocab_file)

        self.special_token_ids = {}
        configured_special_tokens = []
        if special_tokens_file and os.path.exists(special_tokens_file):
            with open(special_tokens_file, "r", encoding="utf-8") as handle:
                metadata = json.load(handle)
            self.special_token_ids = {
                str(token): int(token_id) for token, token_id in metadata.get("special_tokens", {}).items()
            }
            configured_special_tokens = [str(token) for token in metadata.get("configured_special_tokens", [])]
        self.configured_special_tokens = list(configured_special_tokens)
        self.id_to_special_token = {
            int(token_id): str(token) for token, token_id in self.special_token_ids.items()
        }

        merged_additional_tokens = list(additional_special_tokens or [])
        if isinstance(extra_special_tokens, dict):
            merged_additional_tokens.extend(str(token) for token in extra_special_tokens.keys())
        elif isinstance(extra_special_tokens, (list, tuple)):
            merged_additional_tokens.extend(str(token) for token in extra_special_tokens)
        for token in configured_special_tokens:
            if token not in {bos_token, eos_token, unk_token, pad_token} and token not in merged_additional_tokens:
                merged_additional_tokens.append(token)
        merged_additional_tokens = list(dict.fromkeys(merged_additional_tokens))

        super().__init__(
            bos_token=bos_token,
            eos_token=eos_token,
            unk_token=unk_token,
            pad_token=pad_token,
            additional_special_tokens=merged_additional_tokens,
            **kwargs,
        )

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

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

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

    def get_special_token_id(self, token):
        return self.special_token_ids.get(token)

    def _encode_with_registered_special_tokens(self, text):
        if not text:
            return []

        special_tokens = [
            token
            for token in self.configured_special_tokens
            if token and token in text
        ]
        if not special_tokens:
            return list(self.sp_model.encode(text, out_type=int))

        pattern = "(" + "|".join(re.escape(token) for token in sorted(special_tokens, key=len, reverse=True)) + ")"
        parts = re.split(pattern, text)

        encoded = []
        for part in parts:
            if not part:
                continue
            special_token_id = self.special_token_ids.get(part)
            if special_token_id is not None:
                encoded.append(int(special_token_id))
            else:
                encoded.extend(self.sp_model.encode(part, out_type=int))
        return encoded

    def _convert_token_to_id(self, token):
        if token in self.special_token_ids:
            return self.special_token_ids[token]
        piece_id = self.sp_model.piece_to_id(token)
        if piece_id < 0:
            return self.sp_model.unk_id()
        return int(piece_id)

    def _convert_id_to_token(self, index):
        if index in self.id_to_special_token:
            return self.id_to_special_token[index]
        if index in self.added_tokens_decoder:
            return self.added_tokens_decoder[index].content
        return self.sp_model.id_to_piece(int(index))

    def convert_tokens_to_string(self, tokens):
        if not tokens:
            return ""
        return self.sp_model.decode_pieces(tokens)

    def decode(self, token_ids, skip_special_tokens=False, **kwargs):
        del kwargs

        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()

        if isinstance(token_ids, (list, tuple)) and token_ids and isinstance(token_ids[0], (list, tuple)):
            token_ids = token_ids[0]

        if not isinstance(token_ids, list):
            token_ids = [int(token_ids)]
        else:
            token_ids = [int(token_id) for token_id in token_ids]

        if skip_special_tokens:
            special_token_ids = {int(token_id) for token_id in self.special_token_ids.values()}
            token_ids = [token_id for token_id in token_ids if token_id not in special_token_ids]

        return self.sp_model.decode(token_ids)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return list(token_ids_0)
        return list(token_ids_0) + list(token_ids_1)

    def encode(self, text, return_tensors=None, **kwargs):
        import torch

        add_special_tokens = kwargs.pop("add_special_tokens", True)
        del add_special_tokens

        input_ids = self._encode_with_registered_special_tokens(text)
        if return_tensors == "pt":
            return torch.tensor(input_ids, dtype=torch.long).unsqueeze(0)
        return input_ids

    def __call__(self, text, return_tensors=None, **kwargs):
        import torch

        add_special_tokens = kwargs.pop("add_special_tokens", True)
        del add_special_tokens

        is_single = isinstance(text, str)
        texts = [text] if is_single else list(text)
        encoded_batch = [self._encode_with_registered_special_tokens(item) for item in texts]

        padding = kwargs.pop("padding", False)
        truncation = kwargs.pop("truncation", False)
        max_length = kwargs.pop("max_length", None)
        return_attention_mask = kwargs.pop("return_attention_mask", True)

        if truncation and max_length is not None:
            encoded_batch = [ids[:max_length] for ids in encoded_batch]

        if padding == "max_length" and max_length is None:
            raise ValueError("max_length must be set when padding='max_length'")
        if padding:
            target_length = max_length if max_length is not None else max(len(ids) for ids in encoded_batch)
            padded_batch = []
            attention_masks = []
            for ids in encoded_batch:
                trimmed = ids[:target_length]
                pad_len = target_length - len(trimmed)
                padded_batch.append(trimmed + [self.pad_token_id] * pad_len)
                attention_masks.append([1] * len(trimmed) + [0] * pad_len)
        else:
            padded_batch = encoded_batch
            attention_masks = [[1] * len(ids) for ids in encoded_batch]

        if return_tensors == "pt":
            if not padding and len({len(ids) for ids in padded_batch}) > 1:
                raise ValueError("Batch elements must have the same length when return_tensors='pt' without padding")
            input_ids = torch.tensor(padded_batch, dtype=torch.long)
            result = {"input_ids": input_ids}
            if return_attention_mask:
                result["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long)
            if is_single:
                result = {key: value for key, value in result.items()}
            return result

        result = {"input_ids": padded_batch[0] if is_single else padded_batch}
        if return_attention_mask:
            result["attention_mask"] = attention_masks[0] if is_single else attention_masks
        return result

    def build_chat_inputs(self, prompt, return_tensors=None):
        import torch

        user_token_id = self.special_token_ids["<user>"]
        assistant_token_id = self.special_token_ids["<assistant>"]
        prompt_ids = self._encode_with_registered_special_tokens(prompt)
        input_ids = [user_token_id, *prompt_ids, assistant_token_id]
        attention_mask = [1] * len(input_ids)

        if return_tensors == "pt":
            return {
                "input_ids": torch.tensor(input_ids, dtype=torch.long).unsqueeze(0),
                "attention_mask": torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0),
            }

        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
        }

    def save_vocabulary(self, save_directory, filename_prefix=None):
        if not os.path.isdir(save_directory):
            raise ValueError(f"Vocabulary path should be a directory, got: {save_directory}")
        vocab_name = self.vocab_files_names["vocab_file"]
        metadata_name = self.vocab_files_names["special_tokens_file"]
        if filename_prefix:
            vocab_name = f"{filename_prefix}-{vocab_name}"
            metadata_name = f"{filename_prefix}-{metadata_name}"
        out_vocab_file = os.path.join(save_directory, vocab_name)
        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
            shutil.copyfile(self.vocab_file, out_vocab_file)
        outputs = (out_vocab_file,)
        if self.special_tokens_file and os.path.exists(self.special_tokens_file):
            out_metadata_file = os.path.join(save_directory, metadata_name)
            if os.path.abspath(self.special_tokens_file) != os.path.abspath(out_metadata_file):
                shutil.copyfile(self.special_tokens_file, out_metadata_file)
            outputs = outputs + (out_metadata_file,)
        return outputs