"""HuggingFace tokenizer wrapper around tiktoken for Logos."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import tiktoken
from transformers import PreTrainedTokenizer


class LogosTokenizer(PreTrainedTokenizer):
    model_input_names = ["input_ids", "attention_mask"]
    vocab_files_names: Dict[str, str] = {}

    def __init__(
        self,
        encoding_name: str = "cl100k_base",
        errors: str = "replace",
        **kwargs,
    ):
        self.encoding_name = encoding_name
        self.encoding = tiktoken.get_encoding(encoding_name)
        self.errors = errors
        eos = "<|endoftext|>"
        kwargs.setdefault("eos_token", eos)
        kwargs.setdefault("pad_token", eos)
        kwargs.setdefault("unk_token", eos)
        super().__init__(**kwargs)

    @property
    def vocab_size(self) -> int:
        return int(self.encoding.n_vocab)

    def get_vocab(self) -> Dict[str, int]:
        return {str(i): i for i in range(self.vocab_size)}

    def __len__(self) -> int:
        return self.vocab_size

    def _tokenize(self, text: str, **kwargs) -> List[str]:
        ids = self.encoding.encode(
            text,
            allowed_special=kwargs.get("allowed_special", set()),
            disallowed_special=kwargs.get("disallowed_special", ()),
        )
        return [str(i) for i in ids]

    def _convert_token_to_id(self, token: str) -> int:
        if token in {self.eos_token, self.pad_token, self.unk_token}:
            return int(self.encoding.eot_token)
        try:
            return int(token)
        except (TypeError, ValueError):
            return int(self.encoding.eot_token)

    def _convert_id_to_token(self, index: int) -> str:
        if int(index) == int(self.encoding.eot_token):
            return self.eos_token
        return str(int(index))

    def convert_tokens_to_ids(self, tokens):
        if tokens is None:
            return None
        if isinstance(tokens, (list, tuple)):
            return [self._convert_token_to_id(tok) for tok in tokens]
        return self._convert_token_to_id(tokens)

    def convert_ids_to_tokens(self, ids, skip_special_tokens: bool = False):
        if ids is None:
            return None
        if isinstance(ids, (list, tuple)):
            return [self.convert_ids_to_tokens(i, skip_special_tokens=skip_special_tokens) for i in ids]
        idx = int(ids)
        if skip_special_tokens and idx == int(self.encoding.eot_token):
            return None
        return self._convert_id_to_token(idx)

    def convert_tokens_to_string(self, tokens: Iterable[str]) -> str:
        ids = [self._convert_token_to_id(tok) for tok in tokens]
        return self.encoding.decode(ids, errors=self.errors)

    def build_inputs_with_special_tokens(
        self,
        token_ids_0: List[int],
        token_ids_1: Optional[List[int]] = None,
    ) -> List[int]:
        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: List[int],
        token_ids_1: Optional[List[int]] = None,
        already_has_special_tokens: bool = False,
    ) -> List[int]:
        if already_has_special_tokens:
            ids = token_ids_0
        elif token_ids_1 is None:
            ids = token_ids_0
        else:
            ids = token_ids_0 + token_ids_1
        eos_id = int(self.encoding.eot_token)
        return [1 if int(tok) == eos_id else 0 for tok in ids]

    def _decode(
        self,
        token_ids: List[int],
        skip_special_tokens: bool = False,
        clean_up_tokenization_spaces: Optional[bool] = None,
        **kwargs,
    ) -> str:
        ids = [int(i) for i in token_ids]
        if skip_special_tokens:
            eos_id = int(self.encoding.eot_token)
            ids = [i for i in ids if i != eos_id]
        return self.encoding.decode(ids, errors=self.errors)

    def save_vocabulary(
        self,
        save_directory: str,
        filename_prefix: Optional[str] = None,
    ) -> Tuple[str, ...]:
        path = Path(save_directory)
        path.mkdir(parents=True, exist_ok=True)
        name = f"{filename_prefix + '-' if filename_prefix else ''}logos_tokenizer.json"
        out = path / name
        out.write_text(json.dumps({"encoding_name": self.encoding_name}, indent=2))
        return (str(out),)


__all__ = ["LogosTokenizer"]