from __future__ import annotations

from typing import Any

import litellm


def calculate_cost(response: object) -> float:
    """Calculate actual cost of a completion using litellm."""
    try:
        cost = litellm.completion_cost(completion_response=response)
        return float(cost)
    except Exception:
        return 0.0


def estimate_max_cost(
    most_expensive_model: str,
    response: Any = None,
    messages: list[dict] | None = None,
) -> float:
    """Estimate what the most expensive model would have cost.

    If a response is provided, uses its actual token counts for accuracy.
    Otherwise falls back to estimating from message character counts.
    """
    input_cost = get_model_input_cost_per_token(most_expensive_model)
    if input_cost is None:
        return 0.0

    output_cost = _get_model_output_cost_per_token(most_expensive_model) or input_cost * 3

    # Prefer actual token counts from the response
    if response is not None:
        usage = getattr(response, "usage", None)
        if usage is not None:
            input_tokens = getattr(usage, "prompt_tokens", 0) or 0
            output_tokens = getattr(usage, "completion_tokens", 0) or 0
            return (input_tokens * input_cost) + (output_tokens * output_cost)

    # Fall back to character-based estimation
    if messages is not None:
        char_count = sum(len(m.get("content", "")) for m in messages)
        estimated_input_tokens = char_count / 4
        estimated_output_tokens = estimated_input_tokens * 1.5
        return (estimated_input_tokens * input_cost) + (estimated_output_tokens * output_cost)

    return 0.0


def get_model_input_cost_per_token(model: str) -> float | None:
    """Look up per-token input pricing from litellm's database."""
    try:
        model_info = litellm.get_model_info(model)
        if model_info and "input_cost_per_token" in model_info:
            return float(model_info["input_cost_per_token"])
    except Exception:
        pass
    return None


def _get_model_output_cost_per_token(model: str) -> float | None:
    """Look up per-token output pricing from litellm's database."""
    try:
        model_info = litellm.get_model_info(model)
        if model_info and "output_cost_per_token" in model_info:
            return float(model_info["output_cost_per_token"])
    except Exception:
        pass
    return None
