from __future__ import annotations

from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from codeflash_api.llm._models import LLM


def calculate_llm_cost(
    response: Any,
    llm: LLM,
) -> float:
    """
    Calculate cost in USD from a provider response.
    """
    usage = getattr(response, "usage", None)
    if usage is None:
        return 0.0

    input_rate = llm.input_cost or 0.0
    cached_rate = (
        llm.cached_input_cost
        if llm.cached_input_cost is not None
        else input_rate
    )
    output_rate = llm.output_cost or 0.0

    if llm.model_type == "anthropic":
        cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
        cache_creation = getattr(usage, "cache_creation_input_tokens", 0) or 0
        non_cached = usage.input_tokens + cache_creation
        output = usage.output_tokens
    else:
        details = getattr(usage, "prompt_tokens_details", None)
        cache_read = (
            (getattr(details, "cached_tokens", 0) or 0) if details else 0
        )
        non_cached = usage.prompt_tokens - cache_read
        output = usage.completion_tokens

    return float(
        (
            non_cached * input_rate
            + cache_read * cached_rate
            + output * output_rate
        )
        / 1_000_000
    )
