from typing import Literal

import tiktoken


class TokenCounter:
    """A class to count tokens in texts using specified tokenizer models.

    This class supports tokenization using either the `tiktoken` or `transformers` libraries,
    allowing for flexible text processing and token counting.
    """

    def __init__(self, model: str, provider: Literal["tiktoken", "transformers"]):
        """Initializes the TokenCounter with a model and provider.

        Args:
            model (str): The model identifier for the tokenizer.
            provider (Literal["tiktoken", "transformers"]): The library to use for tokenization.
                Must be either 'tiktoken' for tiktoken library or 'transformers' for Hugging Face's transformers
                library. 'transformers' is not yet supported.

        Raises:
            ValueError: If the provider is not 'tiktoken' or 'transformers'.
        """
        self.model = model
        self.provider = provider
        self.texts = []
        self.cnt_tokens_per_text = []
        self._total_tokens = 0
        if provider == "tiktoken":
            self.tokenizer = tiktoken.encoding_for_model(self.model)
        elif provider == "transformers":
            raise NotImplementedError("Transformers library is not yet supported.")
        else:
            raise ValueError("Unrecognized tokenizer library.")

    def add_text(self, text: str) -> int:
        """Adds text to the TokenCounter, updates the token counts and returns them.

        This method encodes the given text using the initialized tokenizer and updates
        the total token count as well as the count of tokens for the individual text.

        Args:
            text (str): The text to be tokenized and counted.

        Returns:
            int: The number of tokens in the added text.
        """
        self.texts.append(text)
        curr_tokens = len(self.tokenizer.encode(text, disallowed_special=()))
        self.cnt_tokens_per_text.append(curr_tokens)
        self._total_tokens += curr_tokens
        return curr_tokens

    @property
    def total_tokens(self) -> int:
        """Retrieves the total number of tokens counted across all added texts.

        Returns:
            int: The total number of tokens.
        """
        return self._total_tokens
