"""
Token counting utilities using tiktoken.

Provides simple token counting using tiktoken for OpenAI-compatible models,
with a heuristic fallback for unsupported models.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import tiktoken as _tiktoken

try:
    import tiktoken

    TIKTOKEN_AVAILABLE = True
except ImportError:
    tiktoken = None  # type: ignore[assignment]
    TIKTOKEN_AVAILABLE = False
    logger.warning("tiktoken not available. Token counting will use estimates.")


def _get_encoding(model_name: str | None) -> _tiktoken.Encoding | None:
    """Get tiktoken encoding for a model, returning None on failure."""
    if not TIKTOKEN_AVAILABLE or tiktoken is None:
        return None

    # Strip provider prefix (e.g. "openai:gpt-4o" -> "gpt-4o") so tiktoken
    # can resolve the correct encoding (o200k_base for GPT-4o+, cl100k_base
    # for GPT-4/3.5-turbo).
    bare_model = model_name or "gpt-4o"
    if ":" in bare_model:
        bare_model = bare_model.split(":", 1)[1]

    try:
        return tiktoken.encoding_for_model(bare_model)
    except KeyError:
        # Model not recognized by tiktoken; fall back to o200k_base which
        # covers GPT-4o, GPT-4.5, GPT-5, o1, o3, and all newer models.
        try:
            return tiktoken.get_encoding("o200k_base")
        except Exception:
            return None


def count_tokens(
    text: str,
    model_name: str | None = None,
    encoding_name: str | None = None,  # noqa: ARG001
) -> int:
    """
    Count tokens in text using tiktoken.

    Args:
        text: Text to count tokens for
        model_name: Model name to get tokenizer for
        encoding_name: Ignored (kept for compatibility)

    Returns:
        Number of tokens, or estimated count if tiktoken unavailable
    """
    encoding = _get_encoding(model_name)
    if encoding is None:
        return estimate_tokens(text)

    try:
        return len(encoding.encode(text))
    except Exception as e:
        logger.warning(f"Failed to count tokens with tiktoken: {e}")
        return estimate_tokens(text)


def estimate_tokens(text: str) -> int:
    """
    Estimate token count using simple heuristics.

    This is a fallback when tiktoken is not available.
    Assumes roughly 4 characters per token on average.

    Args:
        text: Text to estimate tokens for

    Returns:
        Estimated number of tokens
    """
    if not text:
        return 0

    # Simple heuristic: ~4 characters per token for English text
    char_count = len(text)
    word_count = len(text.split())

    # More sophisticated estimate accounting for word boundaries
    estimated = (char_count / 3.5) + (word_count * 0.1)

    return max(1, int(estimated))


def count_message_tokens(messages: list[dict[str, Any]], model_name: str | None = None) -> int:
    """
    Count tokens in a list of messages using tiktoken.

    Accounts for per-message overhead tokens used by the chat format.

    Args:
        messages: List of message dictionaries with 'role' and 'content'
        model_name: Model name for tokenizer selection

    Returns:
        Total token count for all messages
    """
    if not messages:
        return 0

    encoding = _get_encoding(model_name)
    if encoding is None:
        # Fallback: sum individual message tokens using heuristic
        total = 0
        for message in messages:
            content = message.get("content", "")
            if isinstance(content, str):
                total += estimate_tokens(content)
            elif isinstance(content, list):
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "text":
                        total += estimate_tokens(item.get("text", ""))
        return total

    try:
        # Each message has overhead tokens for formatting delimiters.
        # Per OpenAI's reference: 3 overhead tokens per message for GPT-4+.
        # Role and content values are encoded separately below.
        tokens_per_message = 3
        total = 0
        for message in messages:
            total += tokens_per_message
            content = message.get("content", "")
            if isinstance(content, str):
                total += len(encoding.encode(content))
            elif isinstance(content, list):
                for item in content:
                    if isinstance(item, dict) and item.get("type") == "text":
                        total += len(encoding.encode(item.get("text", "")))
            role = message.get("role", "")
            if isinstance(role, str):
                total += len(encoding.encode(role))
        # Every reply is primed with <|start|>assistant<|message|>
        total += 3
        return total
    except Exception as e:
        logger.warning(f"Failed to count message tokens: {e}")
        return sum(
            estimate_tokens(msg.get("content", ""))
            for msg in messages
            if isinstance(msg.get("content"), str)
        )
