"""
Token counting utilities for context management.

Uses tiktoken for token counting. While tiktoken is OpenAI's tokenizer,
cl100k_base encoding provides reasonable estimates for most LLMs.
"""
from functools import lru_cache

from src.shared.logging import get_logger

logger = get_logger(__name__)


@lru_cache(maxsize=1)
def _get_encoding():
    """Get tiktoken encoding (cached)."""
    import tiktoken
    # cl100k_base is used by GPT-4 and provides reasonable estimates for most models
    return tiktoken.get_encoding("cl100k_base")


def count_tokens(text: str) -> int:
    """
    Count tokens in a text string.

    Args:
        text: Text to count tokens for

    Returns:
        Token count
    """
    try:
        encoding = _get_encoding()
        return len(encoding.encode(text))
    except Exception as e:
        # Fallback to rough estimate if tiktoken fails
        logger.warning(f"Token counting failed, using estimate: {e}")
        return len(text) // 4


def count_message_tokens(messages: list[dict[str, str]]) -> int:
    """
    Count tokens for a list of chat messages.

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

    Returns:
        Total token count including message overhead
    """
    try:
        encoding = _get_encoding()
        total = 0
        for msg in messages:
            # Each message has ~4 tokens overhead for role/formatting
            total += 4
            total += len(encoding.encode(msg.get("content", "")))
            total += len(encoding.encode(msg.get("role", "")))
        # Add 2 tokens for assistant response priming
        total += 2
        return total
    except Exception as e:
        # Fallback to rough estimate
        logger.warning(f"Token counting failed, using estimate: {e}")
        total = 0
        for msg in messages:
            total += len(msg.get("content", "")) // 4
            total += 4  # Overhead per message
        return total


def estimate_tokens(text: str) -> int:
    """
    Quick token estimate without external library.

    Uses ~4 characters per token heuristic.
    Less accurate but faster for rough estimates.

    Args:
        text: Text to estimate

    Returns:
        Estimated token count
    """
    return len(text) // 4