"""Token counting utilities for accurate cost estimation.

Uses tiktoken for BPE-based token counting instead of the naive
``len(text) // 4`` heuristic.  For OpenAI models the correct encoding
is selected automatically; for Anthropic/Google models we fall back to
``cl100k_base`` which is a reasonable BPE approximation.
"""

from __future__ import annotations

import functools

import structlog

logger = structlog.get_logger()


@functools.lru_cache(maxsize=16)
def _get_encoding(model: str):
    """Get the tiktoken encoding for a model, with fallback.

    Cached so repeated calls for the same model are free.
    """
    import tiktoken

    try:
        return tiktoken.encoding_for_model(model)
    except KeyError:
        return tiktoken.get_encoding("cl100k_base")


def count_tokens(text: str, model: str = "") -> int:
    """Count tokens in *text* using tiktoken.

    Falls back to ``len(text) // 4`` only if tiktoken is unavailable.
    """
    if not text:
        return 0
    try:
        enc = _get_encoding(model)
        return len(enc.encode(text))
    except Exception:  # defensive-catch: tokenizer fallback to len//4
        return max(1, len(text) // 4)


def count_messages_tokens(
    messages: list,
    model: str = "",
) -> int:
    """Count tokens across a list of chat messages.

    Each message carries ~3 tokens of overhead for role and separators
    (following the OpenAI token counting guide).
    """
    if not messages:
        return 0

    total = 0
    per_message_overhead = 3  # role + separators

    for msg in messages:
        content = msg.content if hasattr(msg, "content") else str(msg)
        total += count_tokens(content, model) + per_message_overhead

    # Every reply is primed with 3 tokens
    total += 3
    return total
