import tiktoken
import json
from typing import Dict, Any

class TokenLimitExceededError(Exception):
    """Custom exception raised when the input payload exceeds the token limit."""

    def __init__(self, message: str, token_count: int, token_limit: int):
        super().__init__(message)
        self.token_count = token_count
        self.token_limit = token_limit



_CACHED_ENCODINGS = {}

class Validator:
    def __init__(self, token_limit: int = 100000, encoding_name: str = "cl100k_base"):
        """
        Initializes the Validator with a token limit and encoding.

        Args:
            token_limit (int): The maximum number of tokens allowed.
            encoding_name (str): The name of the tokenizer encoding to use.
        """
        self.token_limit = token_limit
        
        if encoding_name not in _CACHED_ENCODINGS:
            try:
                _CACHED_ENCODINGS[encoding_name] = tiktoken.get_encoding(encoding_name)
            except Exception as e:
                import logging
                logging.error(f"[Validator] Failed to initialize tiktoken encoding '{encoding_name}': {e}")
                # Fallback to None if internet not available to download vocab files
                _CACHED_ENCODINGS[encoding_name] = None
        
        self.encoding = _CACHED_ENCODINGS.get(encoding_name)

    def get_token_count(self, text: str) -> int:
        """Helper to get token count with fallback if tiktoken is unavailable."""
        if not text:
            return 0
        if self.encoding:
            try:
                return len(self.encoding.encode(text))
            except Exception:
                pass
        # Rough estimate: 3.5 chars per token for English
        return len(text) // 4

    def precheck_tokensize(self, input_payload: Dict[str, Any]) -> None:
        """
        Checks if the input payload's token count exceeds the configured limit.

        Args:
            input_payload (Dict[str, Any]): The payload to be checked.

        Raises:
            TokenLimitExceededError: If the payload's token count is too high.
        """
        payload_string: str = json.dumps(input_payload)
        token_count: int = self.get_token_count(payload_string)
        
        if token_count > self.token_limit:
            raise TokenLimitExceededError(
                f"Input payload token count ({token_count}) exceeds the limit of {self.token_limit} tokens.",
                token_count,
                self.token_limit,
            )