from __future__ import annotations

import json
import os

from transformers import PreTrainedTokenizer


class TransWaveByteTokenizer(PreTrainedTokenizer):

    vocab_files_names = {
        "vocab_file": "vocab.json"
    }

    model_input_names = [
        "input_ids"
    ]


    def __init__(
        self,
        vocab_file=None,
        model_max_length=20480,
        **kwargs,
    ):

        self.vocab_file = vocab_file

        self._id_to_token = [
            f"<0x{i:02X}>"
            for i in range(256)
        ]

        self._token_to_id = {
            token: i
            for i, token
            in enumerate(self._id_to_token)
        }

        # Values may already arrive through tokenizer_config.json.
        # Never pass the same keyword twice.

        kwargs.setdefault(
            "model_max_length",
            model_max_length,
        )

        kwargs.setdefault(
            "clean_up_tokenization_spaces",
            False,
        )

        super().__init__(**kwargs)


    @property
    def vocab_size(self):

        return 256


    def get_vocab(self):

        vocab = dict(
            self._token_to_id
        )

        vocab.update(
            self.added_tokens_encoder
        )

        return vocab


    def _tokenize(
        self,
        text,
        **kwargs,
    ):

        return [
            self._id_to_token[b]
            for b in text.encode("utf-8")
        ]


    def _convert_token_to_id(
        self,
        token,
    ):

        if token not in self._token_to_id:
            raise ValueError(
                f"Unknown byte token: {token}"
            )

        return self._token_to_id[token]


    def _convert_id_to_token(
        self,
        index,
    ):

        index = int(index)

        if not 0 <= index < 256:
            raise ValueError(
                f"Byte id out of range: {index}"
            )

        return self._id_to_token[index]


    def convert_tokens_to_string(
        self,
        tokens,
    ):

        data = bytearray()

        for token in tokens:

            if token in self._token_to_id:
                data.append(
                    self._token_to_id[token]
                )

        return bytes(data).decode(
            "utf-8",
            errors="replace",
        )


    def save_vocabulary(
        self,
        save_directory,
        filename_prefix=None,
    ):

        os.makedirs(
            save_directory,
            exist_ok=True,
        )

        filename = (
            (
                filename_prefix + "-"
            )
            if filename_prefix
            else ""
        ) + "vocab.json"

        path = os.path.join(
            save_directory,
            filename,
        )

        vocab = {
            token: i
            for i, token
            in enumerate(self._id_to_token)
        }

        with open(
            path,
            "w",
            encoding="utf-8",
        ) as f:

            json.dump(
                vocab,
                f,
                ensure_ascii=False,
                indent=2,
            )

        return (path,)