"""Author RAG Chatbot SaaS — Token Counter Utility.

Counts tokens for gpt-4o using tiktoken before sending to OpenAI.
RULE: Always call count_tokens() before assembling the context
to ensure we stay within RAG_MAX_CONTEXT_TOKENS.
"""

import structlog
import tiktoken

logger = structlog.get_logger(__name__)

# gpt-4o uses the same tokenizer as gpt-4
_ENCODING_NAME = "cl100k_base"
_encoding: tiktoken.Encoding | None = None


def _get_encoding() -> tiktoken.Encoding:
    """Lazily load and cache the tiktoken encoding."""
    global _encoding
    if _encoding is None:
        _encoding = tiktoken.get_encoding(_ENCODING_NAME)
    return _encoding


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

    Args:
        text: Input string to tokenize.

    Returns:
        Integer token count.
    """
    return len(_get_encoding().encode(text))


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

    Accounts for per-message overhead (role + separators).

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

    Returns:
        Total token count including overhead.
    """
    encoding = _get_encoding()
    total = 0
    for msg in messages:
        # Each message has: 4 overhead tokens + role + content
        total += 4
        total += len(encoding.encode(msg.get("role", "")))
        total += len(encoding.encode(msg.get("content", "")))
    total += 2  # Reply priming
    return total


def trim_messages_to_budget(
    messages: list[dict],
    max_tokens: int,
    preserve_system: bool = True,
) -> list[dict]:
    """Trim conversation history to fit within a token budget.

    Removes oldest messages first. System message is always preserved
    if preserve_system=True.

    Args:
        messages: Full list of chat messages (oldest first).
        max_tokens: Maximum allowed token count.
        preserve_system: If True, never remove the system message.

    Returns:
        Trimmed list of messages that fits within max_tokens.
    """
    if count_messages_tokens(messages) <= max_tokens:
        return messages

    system_msgs = [m for m in messages if m["role"] == "system"] if preserve_system else []
    history_msgs = [m for m in messages if m["role"] != "system"]

    while history_msgs and count_messages_tokens(system_msgs + history_msgs) > max_tokens:
        history_msgs.pop(0)  # Remove oldest non-system message
        logger.debug("Trimmed oldest message from context", remaining=len(history_msgs))

    return system_msgs + history_msgs


def fits_budget(text: str, budget: int) -> bool:
    """Check whether a text fits within a token budget.

    Args:
        text: Input string to evaluate.
        budget: Maximum allowed token count.

    Returns:
        True if token count is within budget, False if it exceeds it.
    """
    return count_tokens(text) <= budget