import json
from pathlib import Path
from typing import List, Optional

from transformers import PreTrainedTokenizer


class KsByteTokenizer(PreTrainedTokenizer):
    """UTF-8 byte tokenizer for ks_byte_lm.

    IDs 0..255 are raw UTF-8 byte values; 256/257/258 are BOS/EOS/PAD.
    """

    model_input_names = ["input_ids", "attention_mask"]

    def __init__(
        self,
        vocab_file: Optional[str] = None,
        zwnj_policy="keep",
        digit_policy="keep_native",
        remove_diacritics=False,
        **kwargs,
    ):
        self.byte_vocab = 256
        self.vocab = {f"byte_{i}": i for i in range(256)}
        self.vocab.update({"<bos>": 256, "<eos>": 257, "<pad>": 258})
        self.ids_to_tokens = {v: k for k, v in self.vocab.items()}
        self.zwnj_policy = zwnj_policy
        self.digit_policy = digit_policy
        self.remove_diacritics = remove_diacritics
        kwargs.pop("bos_token", None)
        kwargs.pop("eos_token", None)
        kwargs.pop("pad_token", None)
        kwargs.pop("unk_token", None)
        kwargs.pop("model_max_length", None)
        super().__init__(
            bos_token="<bos>",
            eos_token="<eos>",
            pad_token="<pad>",
            unk_token="<pad>",
            model_max_length=2048,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return 259

    def get_vocab(self):
        return dict(self.vocab)

    def _tokenize(self, text: str) -> List[str]:
        return [f"byte_{b}" for b in text.encode("utf-8")]

    def _convert_token_to_id(self, token):
        return self.vocab.get(token, 258)

    def _convert_id_to_token(self, index):
        return self.ids_to_tokens.get(int(index), "<pad>")

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

    def save_vocabulary(self, save_directory, filename_prefix=None):
        path = Path(save_directory) / ((filename_prefix + "-" if filename_prefix else "") + "vocab.json")
        path.write_text(json.dumps(self.vocab, ensure_ascii=False, indent=2), encoding="utf-8")
        return (str(path),)

    def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=None, **kwargs):
        if hasattr(token_ids, "tolist"):
            token_ids = token_ids.tolist()
        if token_ids and isinstance(token_ids[0], list):
            token_ids = token_ids[0]
        bs = []
        for i in token_ids:
            i = int(i)
            if 0 <= i < 256:
                bs.append(i)
            elif not skip_special_tokens and i in (256, 257, 258):
                continue
        return bytes(bs).decode("utf-8", errors="replace")