from __future__ import annotations

import json
from pathlib import Path

from transformers import PreTrainedTokenizer


class LaGrangeFlashTokenizer(PreTrainedTokenizer):
    """Lossless UTF-8 byte tokenizer for the interoperability bridge."""

    vocab_files_names = {"vocab_file": "vocab.json"}
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, vocab_file=None, **kwargs):
        self.vocab_file = vocab_file
        self._special = {"<pad>": 0, "<bos>": 1, "<eos>": 2, "<unk>": 3}
        super().__init__(
            pad_token="<pad>",
            bos_token="<bos>",
            eos_token="<eos>",
            unk_token="<unk>",
            **kwargs,
        )

    @property
    def vocab_size(self):
        return 260

    def get_vocab(self):
        vocab = dict(self._special)
        vocab.update({f"<0x{i:02X}>": i + 4 for i in range(256)})
        return vocab

    def _tokenize(self, text, **kwargs):
        return [f"<0x{byte:02X}>" for byte in str(text).encode("utf-8")]

    def _convert_token_to_id(self, token):
        if token in self._special:
            return self._special[token]
        if token.startswith("<0x") and token.endswith(">"):
            try:
                value = int(token[3:-1], 16)
            except ValueError:
                return self._special["<unk>"]
            if 0 <= value <= 255:
                return value + 4
        return self._special["<unk>"]

    def _convert_id_to_token(self, index):
        special_by_id = {value: key for key, value in self._special.items()}
        if index in special_by_id:
            return special_by_id[index]
        if 4 <= index < 260:
            return f"<0x{index - 4:02X}>"
        return "<unk>"

    def convert_tokens_to_string(self, tokens):
        data = bytearray()
        parts = []

        def flush_bytes():
            nonlocal data
            if data:
                parts.append(bytes(data).decode("utf-8", errors="replace"))
                data = bytearray()

        for token in tokens:
            if token.startswith("<0x") and token.endswith(">"):
                try:
                    data.append(int(token[3:-1], 16))
                    continue
                except ValueError:
                    pass
            flush_bytes()
            if token not in self._special:
                parts.append(token)

        flush_bytes()
        return "".join(parts)

    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
        ids = [self.bos_token_id] + list(token_ids_0) + [self.eos_token_id]
        if token_ids_1 is not None:
            ids += list(token_ids_1) + [self.eos_token_id]
        return ids

    def get_special_tokens_mask(
        self,
        token_ids_0,
        token_ids_1=None,
        already_has_special_tokens=False,
    ):
        if already_has_special_tokens:
            special_ids = {
                self.pad_token_id,
                self.bos_token_id,
                self.eos_token_id,
                self.unk_token_id,
            }
            return [1 if token_id in special_ids else 0 for token_id in token_ids_0]

        mask = [1] + [0] * len(token_ids_0) + [1]
        if token_ids_1 is not None:
            mask += [0] * len(token_ids_1) + [1]
        return mask

    def save_vocabulary(self, save_directory, filename_prefix=None):
        directory = Path(save_directory)
        directory.mkdir(parents=True, exist_ok=True)
        prefix = f"{filename_prefix}-" if filename_prefix else ""
        path = directory / f"{prefix}vocab.json"
        path.write_text(json.dumps(self.get_vocab(), indent=2), encoding="ascii")
        return (str(path),)