import re
from .initialize import init_models_config
from .providers.factory import get_provider

class TokenLimits:
    """
    Define LLM model parameters and token count of messages.
    First-time execution creates a models.yaml file with example models,
    which can be expanded by the user for additional inputs or modifications like tokens per message.
    Token counting is delegated to each provider: OpenAI uses tiktoken, Anthropic uses the SDK.

    Sources:
    - https://platform.openai.com/docs/models
    - https://docs.anthropic.com/en/docs/about-claude/models
    """
    def __init__(self, model="gpt-4o-mini", yaml_file="models.yaml"):
        # The model name may also be an OpenAI fine-tuned model, which has the base model
        # after "ft:" up to the next colon (:), like "ft:gpt-4o-mini:..."
        self.model = re.search("^ft:([^:]*)", model).group(1) if model.startswith("ft:") else model
        self.config = init_models_config(yaml_file)

        # Get model parameters by configuration file for the:
        # > Maximum number of tokens
        # > Tokens per message
        # > Tokens per name by model
        self.param = None
        for key, param in self.config.items():
            # Set parameters if the configuration key matches either:
            #  > Full model name (an exact match to stop searching)
            #  > Base model name after "ft:" up to next colon (:), for fine-tuned OpenAI models
            if key == model or key == self.model:
                self.param = param
                if key == model:
                    break

        # If the parameters are not set, the model name is invalid or
        # has an undefined model configuration not set in the YAML file
        if self.param is None:
            raise ValueError(
                f"Model \"{model}\" is invalid or its base model is not in: {yaml_file}")

        # Set token model parameter defaults, unless defined in configuration:
        if 'max_tokens' not in self.param:
            self.param['max_tokens'] = 4097

    def max_tokens(self) -> int:
        """Query the Generative AI model's maximum amount of tokens."""
        return self.param['max_tokens']

    async def num_tokens_from_messages(self, messages: list[dict]) -> int:
        """
        Return the number of tokens in messages for this TokenLimits's model.

        Delegates token counting to the appropriate provider:
        - OpenAI models (gpt-*): Uses tiktoken encoding directly
        - Anthropic models (claude-*): Uses the Anthropic SDK beta token counting API

        Args:
            messages: List of message dictionaries with 'role' and 'content' keys.

        Returns:
            The total token count for the messages.

        Raises:
            ProviderAuthError: If API credentials are missing or invalid.
            ProviderConnectionError: If the provider API cannot be reached.
            Exception: If token counting fails for any other reason.

        Sources:
            - https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
            - https://docs.anthropic.com/en/docs/build-a-bot/count-tokens
        """
        return await get_provider(self.model).count_tokens(self.model, messages)
