from typing import Generator
from chunkipy.size_estimators.base_size_estimator import BaseSizeEstimator
from chunkipy.utils import import_dependencies


class OpenAISizeEstimator(BaseSizeEstimator):
    """Estimate size using a ``tiktoken`` encoding compatible with OpenAI models."""

    def __init__(self, encoding: str = "cl100k_base"):
        """Initialize the estimator with the encoding name to load via ``tiktoken``."""
        super().__init__()
        tiktoken = import_dependencies(extra="tiktoken", package_name="tiktoken")
        self.tokenizer = tiktoken.get_encoding(encoding)

    def estimate_size(self, text: str) -> int:
        """
        Estimate the size of the given text using OpenAI's tokenization.

        Args:
            text (str): The text to estimate the size of.

        Returns:
            int: The estimated size of the text in tokens.
        """
        return len(self.tokenizer.encode(text))
    
    def segment(self, text: str) -> Generator[str, None, None]:
        """
        Generate token segments from the given text using OpenAI's tokenization.
        Args:
            text (str): The text to segment.

        Yields:
            str: A single token as segmented by the tokenizer.
        """
        encoded_tokens = self.tokenizer.encode(text)
        for token in (encoded_tokens):
            yield self.tokenizer.decode([token])