import logging
from typing import Any, overload

import tiktoken
from litellm import cost_per_token, token_counter
from pydantic import BaseModel

logger = logging.getLogger(__name__)

TIKTOKEN_ENCODER = "cl100k_base"


class TokenCountCost(BaseModel):
    """Data structure for token counts and costs."""

    prompt_token_cost: float | None = None
    completion_token_cost: float | None = None
    total_token_cost: float | None = None
    prompt_token_count: int | None = None
    completion_token_count: int | None = None
    total_token_count: int | None = None


class TokenCounter:
    # Chunk size for processing long texts (words)
    CHUNK_SIZE = 1000

    def __init__(self, model: str = TIKTOKEN_ENCODER):
        """Initializes a titoken encoder

        :param model: tiktoken model encoder
        """
        self.encoder = tiktoken.get_encoding(model)

    def count(self, query: str | None) -> int:
        """Returns token count of the query using chunking for long texts.

        :param query: string query sent to LLM
        """
        if not query:
            return 0

        # Split into words
        words = query.split()

        # For short texts, encode directly
        if len(words) <= self.CHUNK_SIZE:
            return len(self.encoder.encode(query))

        # For long texts, process in word-based chunks
        total_tokens = 0
        for i in range(0, len(words), self.CHUNK_SIZE):
            chunk_words = words[i : i + self.CHUNK_SIZE]
            chunk_text = " ".join(chunk_words)
            total_tokens += len(self.encoder.encode(chunk_text))

        return total_tokens


def compute_cost_from_tokens(
    model_name: str,
    input_tokens: int = 0,
    output_tokens: int = 0,
) -> float | None:
    """
    Compute cost from token counts.

    Args:
        model_name: The model name (e.g., "gpt-4", "claude-3-opus")
        input_tokens: Number of input tokens
        output_tokens: Number of output tokens

    Returns:
        Cost or None if computation fails
    """
    try:
        if input_tokens == 0 and output_tokens == 0:
            return None

        # Calculate costs using litellm's cost_per_token
        prompt_cost, completion_cost = cost_per_token(
            model=model_name,
            prompt_tokens=input_tokens,
            completion_tokens=output_tokens,
        )

        total_cost = float(prompt_cost) + float(completion_cost)
        return total_cost

    except Exception as e:
        logger.warning(
            f"Error computing cost for model {model_name}: {e}",
        )
        return None


def count_tokens_from_string(
    text: str,
    model_name: str | None = None,
) -> int | None:
    """Calculate token count from string using litellm or fallback to tiktoken."""
    if not text:
        return None
    try:
        if model_name:
            return int(token_counter(model=model_name, text=text))
        else:
            # Fallback: use default tiktoken encoder
            counter = TokenCounter(TIKTOKEN_ENCODER)
            return counter.count(text)
    except Exception as e:
        logger.warning(f"Error counting tokens from string: {e}")
        return None


def count_tokens_from_messages(
    messages: list[dict[str, Any]],
    model_name: str | None = None,
) -> int | None:
    """Calculate token count from messages using litellm or fallback to tiktoken.

    Expects OpenInference normalized format where each message has nested structure:
    {"message": {"role": "user", "content": "..."}}
    """
    if not messages:
        return None

    try:
        formatted = []
        for msg in messages:
            if not isinstance(msg, dict):
                continue

            message_data = msg.get("message", {})
            role = message_data.get("role")
            content = message_data.get("content")

            # Only include messages with content
            if content:
                formatted.append(
                    {"role": str(role) if role else "user", "content": str(content)},
                )

        if not formatted:
            return None

        if model_name:
            # Use litellm for model-specific counting
            return int(token_counter(model=model_name, messages=formatted))
        else:
            # Fallback: concatenate message content and use string counter
            combined_text = "\n".join(
                f"{msg['role']}: {msg['content']}" for msg in formatted
            )
            return count_tokens_from_string(combined_text, model_name=None)
    except Exception as e:
        logger.warning(f"Error counting tokens from messages: {e}")
        return None


@overload
def safe_add(current: None, value: None) -> None: ...


@overload
def safe_add(current: int, value: int | None) -> int: ...


@overload
def safe_add(current: int | None, value: int) -> int: ...


@overload
def safe_add(current: float, value: int | float | None) -> float: ...


@overload
def safe_add(current: int | float | None, value: float) -> float: ...


def safe_add(
    current: int | float | None,
    value: int | float | None,
) -> int | float | None:
    """
    NULL-safe addition for numeric values (token counts or costs).

    Returns None if both values are None.
    Returns the sum if at least one value is not None (treating None as 0).

    Args:
        current: Current accumulated value (or None)
        value: Value to add (or None)

    Returns:
        Sum of values, or None if both inputs are None

    Examples:
        safe_add(None, None) -> None
        safe_add(100, None) -> 100
        safe_add(None, 50) -> 50
        safe_add(100, 50) -> 150
        safe_add(1.5, None) -> 1.5
        safe_add(None, 0.75) -> 0.75
        safe_add(1.5, 0.75) -> 2.25
    """
    if current is None and value is None:
        return None
    return (current or 0) + (value or 0)
