"""A byte-level tokenizer for the private LLM wrapper.

This tokenizer exists so standard Hugging Face text-generation pipelines can
encode a prompt before the wrapper calls the private Python script.
"""

import json
from pathlib import Path

from transformers import PreTrainedTokenizer


class PrivateLLMTokenizer(PreTrainedTokenizer):
    """UTF-8 byte tokenizer with a tiny fixed vocabulary."""

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

    def __init__(
        self,
        vocab_file=None,
        unk_token="<unk>",
        pad_token="<pad>",
        bos_token="<bos>",
        eos_token="<eos>",
        **kwargs,
    ):
        self.vocab_file = vocab_file
        self.byte_tokens = {f"<0x{idx:02X}>": idx for idx in range(256)}
        self.special_tokens = {
            pad_token: 256,
            eos_token: 257,
            bos_token: 258,
            unk_token: 259,
        }
        self.ids_to_tokens = {
            idx: token for token, idx in {**self.byte_tokens, **self.special_tokens}.items()
        }
        super().__init__(
            unk_token=unk_token,
            pad_token=pad_token,
            bos_token=bos_token,
            eos_token=eos_token,
            **kwargs,
        )

    @property
    def vocab_size(self):
        return 260

    def get_vocab(self):
        return {**self.byte_tokens, **self.special_tokens}

    def _tokenize(self, text):
        if not isinstance(text, str):
            text = json.dumps(text, default=str)
        return [f"<0x{byte:02X}>" for byte in text.encode("utf-8")]

    def _convert_token_to_id(self, token):
        return self.get_vocab().get(token, self.unk_token_id)

    def _convert_id_to_token(self, index):
        return self.ids_to_tokens.get(int(index), self.unk_token)

    def convert_tokens_to_string(self, tokens):
        byte_values = []
        for token in tokens:
            if token.startswith("<0x") and token.endswith(">"):
                try:
                    byte_values.append(int(token[3:-1], 16))
                except ValueError:
                    continue
        return bytes(byte_values).decode("utf-8", errors="replace")

    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 get_special_tokens_mask(self, token_ids_0, token_ids_1=None, already_has_special_tokens=False):
        if already_has_special_tokens:
            return [1 if token_id >= 256 else 0 for token_id in token_ids_0]
        if token_ids_1 is None:
            return [0] * len(token_ids_0)
        return [0] * (len(token_ids_0) + len(token_ids_1))

    def create_token_type_ids_from_sequences(self, token_ids_0, token_ids_1=None):
        if token_ids_1 is None:
            return [0] * len(token_ids_0)
        return [0] * (len(token_ids_0) + len(token_ids_1))

    def save_vocabulary(self, save_directory, filename_prefix=None):
        path = Path(save_directory)
        path.mkdir(parents=True, exist_ok=True)
        name = "private_llm_tokenizer.json"
        if filename_prefix:
            name = f"{filename_prefix}-{name}"
        output_path = path / name
        payload = {
            "type": "utf8-byte-tokenizer",
            "vocab_size": self.vocab_size,
            "special_tokens": self.special_tokens,
        }
        output_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
        return (str(output_path),)