"""
Token accounting for the kvict Proxy.

Counts tokens, computes savings (baseline vs actual), and converts to
the hero metrics: "tokens saved," "extra messages gained," "$ saved."

For OpenAI models, uses tiktoken for exact counts.
For other providers, falls back to a character-based heuristic.
"""

from __future__ import annotations

import logging
from typing import List

from proxy.schemas import ChatMessage, TokenPricing, TokenSavings

logger = logging.getLogger(__name__)

# Lazy-loaded tiktoken encoding
_encoding = None


def _get_encoding(model: str):
    """Get tiktoken encoding for a model, falling back to cl100k_base."""
    global _encoding
    try:
        import tiktoken
    except ImportError:
        return None

    if _encoding is not None:
        return _encoding

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


def count_tokens(messages: List[ChatMessage], model: str) -> int:
    """Count prompt tokens for a list of messages.

    Uses tiktoken for OpenAI-family models; falls back to character heuristic.

    The OpenAI token counting formula accounts for per-message overhead
    (role, separators) as documented by OpenAI's cookbook.

    Args:
        messages: List of chat messages.
        model: Model identifier.

    Returns:
        Estimated prompt token count.
    """
    enc = _get_encoding(model)

    if enc is not None:
        # tiktoken-based counting (OpenAI cookbook formula)
        tokens_per_message = 3  # <|start|>{role}\n ... \n
        tokens_per_name = 1
        total = 0
        for msg in messages:
            total += tokens_per_message
            total += len(enc.encode(msg.role))
            if msg.content:
                total += len(enc.encode(msg.content))
            if msg.name:
                total += tokens_per_name
                total += len(enc.encode(msg.name))
        total += 3  # assistant priming
        return total

    # Fallback: ~4 chars per token (rough but reasonable)
    chars = sum(
        len(msg.role) + len(msg.content or "") + len(msg.name or "") for msg in messages
    )
    return max(1, chars // 4)


def compute_savings(
    baseline_tokens: int,
    actual_tokens: int,
    completion_tokens: int,
    pricing: TokenPricing,
    cache_hit: bool = False,
) -> TokenSavings:
    """Compute the savings from proxy optimization.

    Args:
        baseline_tokens: Prompt tokens that WOULD have been sent without proxy.
        actual_tokens: Prompt tokens actually sent (0 if full cache hit).
        completion_tokens: Completion tokens generated.
        pricing: Per-model token pricing.
        cache_hit: Whether this was a full cache hit.

    Returns:
        TokenSavings with all computed fields.
    """
    cached_tokens = baseline_tokens - actual_tokens

    if cache_hit:
        # Full cache hit: entire request was free (no API call at all)
        baseline_cost = (
            baseline_tokens * pricing.input_cost_per_1k / 1000
            + completion_tokens * pricing.output_cost_per_1k / 1000
        )
        actual_cost = 0.0
    else:
        baseline_cost = baseline_tokens * pricing.input_cost_per_1k / 1000
        actual_cost = actual_tokens * pricing.input_cost_per_1k / 1000
        # Completion cost is the same either way (not cached)

    dollars_saved = max(0.0, baseline_cost - actual_cost)

    return TokenSavings(
        baseline_tokens=(
            baseline_tokens + completion_tokens if cache_hit else baseline_tokens
        ),
        actual_tokens=actual_tokens,
        cached_tokens=cached_tokens + completion_tokens if cache_hit else cached_tokens,
        dollars_saved=round(dollars_saved, 6),
        cache_hit=cache_hit,
    )


def tokens_to_messages(tokens_saved: int, avg_per_msg: int = 750) -> int:
    """Convert saved tokens to 'extra messages gained.'

    This is the hero metric for subscription stretching:
    "You gained 12 extra messages this session."

    Args:
        tokens_saved: Total tokens saved by proxy caching.
        avg_per_msg: Average tokens per message (configurable, default 750).

    Returns:
        Number of effective extra messages.
    """
    if avg_per_msg <= 0:
        return 0
    return tokens_saved // avg_per_msg
