import json

from pathlib import Path

from typing import Optional, Union

import torch

class Tokenizer:

    def __init__(self, checkpoint_dir: Union[Path, str]) -> None:

        checkpoint_dir = Path(checkpoint_dir)

        if not checkpoint_dir.exists():

            raise NotADirectoryError(f"The checkpoint directory does not exist: {str(checkpoint_dir)}")

        self.use_bos = self.check_if_bos_token_used(checkpoint_dir)

        self.bos_id = None

        self.eos_id = None

        # some checkpoints have both files, `.model` takes precedence

        if (vocabulary_path := checkpoint_dir / "tokenizer.model").is_file():

            from sentencepiece import SentencePieceProcessor

            self.processor = SentencePieceProcessor(model_file=str(vocabulary_path))

            self.backend = "sentencepiece"

            self.bos_id = self.processor.bos_id()

            self.eos_id = self.processor.eos_id()

        elif (vocabulary_path := checkpoint_dir / "tokenizer.json").is_file():

            from tokenizers import Tokenizer as HFTokenizer

            self.processor = HFTokenizer.from_file(str(vocabulary_path))

            self.backend = "huggingface"

            if (special_tokens_path := checkpoint_dir / "tokenizer_config.json").is_file():

                with open(special_tokens_path) as fp:

                    config = json.load(fp)

                bos_token = config.get("bos_token")

                self.bos_id = self.token_to_id(bos_token) if bos_token is not None else None

                eos_token = config.get("eos_token")

                self.eos_id = self.token_to_id(eos_token) if eos_token is not None else None

            if (special_tokens_path := checkpoint_dir / "generation_config.json").is_file():

                with open(special_tokens_path) as fp:

                    config = json.load(fp)

                if self.bos_id is None:

                    self.bos_id = config.get("bos_token_id")

                if self.eos_id is None:

                    self.eos_id = config.get("eos_token_id")

        else:

            raise NotImplementedError

    @property

    def vocab_size(self) -> int:

        if self.backend == "huggingface":

            return self.processor.get_vocab_size(with_added_tokens=False)

        if self.backend == "sentencepiece":

            return self.processor.vocab_size()

        raise RuntimeError

    def token_to_id(self, token: str) -> int:

        if self.backend == "huggingface":

            id_ = self.processor.token_to_id(token)

        elif self.backend == "sentencepiece":

            id_ = self.processor.piece_to_id(token)

        else:

            raise RuntimeError

        if id_ is None:

            raise ValueError(f"token {token!r} not found in the collection.")

        return id_

    def check_if_bos_token_used(self, checkpoint_dir: Path) -> bool:

        if not (tokenizer_config_path := checkpoint_dir / "tokenizer_config.json").is_file():

            return False

        with open(tokenizer_config_path) as fp:

            config = json.load(fp)

        if any(config.get(check, False) for check in ("add_bos_token", "add_prefix_space")):

            return True

        # for examples that also use the Llama tokenizer, but do not have or set add_bos_token to True.

        # ex: https://huggingface.co/stabilityai/StableBeluga2/blob/main/tokenizer_config.json#L2

        return config.get("add_bos_token") is None and config.get("tokenizer_class") == "LlamaTokenizer"

    def encode(

        self,

        string: str,

        device: Optional[torch.device] = None,

        bos: Optional[bool] = None,

        eos: bool = False,

        max_length: int = -1,

    ) -> torch.Tensor:

        if self.backend == "huggingface":

            tokens = self.processor.encode(string).ids

        elif self.backend == "sentencepiece":

            tokens = self.processor.encode(string)

        else:

            raise RuntimeError

        if bos or (bos is None and self.use_bos):

            bos_id = self.bos_id

            if bos_id is None:

                raise NotImplementedError("This tokenizer does not have a defined a bos token")

            tokens = [bos_id] + tokens

        if eos:

            tokens = tokens + [self.eos_id]

        if max_length > 0:

            tokens = tokens[:max_length]

        return torch.tensor(tokens, dtype=torch.int, device=device)

    def decode(self, tensor: torch.Tensor) -> str:

        tokens = [tensor.item()] if tensor.ndim == 0 else tensor.tolist()

        return self.processor.decode(tokens)