import re
import tiktoken
from .utils import read_yaml
from .initialize import init_tokenGPT_config

class TokenGPT:
    """
    Define LLM model parameters and token count of messages.
    First-time execution creates a tokenGPT.yaml file with example models,
    which can be expanded by the user for additional inputs or modifications like tokens per message.
    OpenAI models use tiktoken for counting; other models use a character-based approximation.

    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="tokenGPT.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_path = init_tokenGPT_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 read_yaml(self.config_path).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:\n{self.config_path}")

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

        # Claude models are not supported by tiktoken; use a character-based approximation
        # (~4 chars per token) instead. A future version may add SDK-native Anthropic counting.
        self._use_char_approximation = self.model.startswith('claude-')
        self.encoding = None
        if not self._use_char_approximation:
            try:
                self.encoding = tiktoken.encoding_for_model(self.model)
            except Exception:
                fallback = "o200k_base" if re.search("^gpt-[45]", self.model) else "cl100k_base"
                try:
                    self.encoding = tiktoken.get_encoding(fallback)
                except Exception:
                    self.encoding = tiktoken.get_encoding("cl100k_base")
                print(f"Warning: using {self.encoding.name} encoding for model \"{self.model}\"")

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

    def num_tokens_from_messages(self, messages: dict[str, str]) -> int:
        """
        Return the number of tokens based on a list of messages and the model in use.

        For OpenAI models: uses tiktoken for an accurate count.
        For Claude models: uses a character-based approximation (~4 chars per token).

        Source: https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
        """
        if self._use_char_approximation:
            return sum(max(1, len(v) // 4) for msg in messages for v in msg.values())

        num_tokens = 0
        for message in messages:
            num_tokens += self.param['tokens_per_message']
            for key, value in message.items():
                num_tokens += len(self.encoding.encode(value))
                if key == "name":
                    num_tokens += self.param['tokens_per_name']
        num_tokens += 3  # Every reply is primed with <|start|>assistant<|message|>
        return num_tokens
