"""
Cost calculation utilities for LLM usage.

Uses genai-prices (pydantic-ai's pricing library) for cache-aware pricing.
Cache read tokens are billed at a reduced rate (~10% for Anthropic, ~50% for OpenAI)
and cache creation tokens at a premium rate (~125% for Anthropic).
"""
from __future__ import annotations

import logging
from typing import Any

logger = logging.getLogger(__name__)

try:
    from genai_prices import Usage as GenAIUsage
    from genai_prices import calc_price as genai_calc_price

    GENAI_PRICES_AVAILABLE = True
except ImportError:
    GENAI_PRICES_AVAILABLE = False
    logger.warning("genai-prices not available. Cost calculation will return 0.")


def _safe_int(value: Any) -> int:  # noqa: ANN401
    """Safely convert a value to int, handling Mock objects and None."""
    if value is None:
        return 0
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def strip_provider_prefix(model_name: str) -> str:
    """Strip provider prefix from model name (e.g. 'openai:gpt-4o' -> 'gpt-4o')."""
    if ":" in model_name:
        return model_name.split(":", 1)[1]
    return model_name


def infer_provider(model_name: str, default: str | None = None) -> str | None:
    """Infer provider ID from model name if it contains a provider prefix.

    Args:
        model_name: Model name, optionally prefixed (e.g. "openai:gpt-4o").
        default: Fallback provider when no prefix is present.
            Use None to let genai-prices auto-detect.
            Use "openai" when the caller knows only OpenAI models are used.

    Returns:
        Provider string or *default* when no prefix is present.
    """
    if ":" in model_name:
        return model_name.split(":", 1)[0]
    return default


def _normalize_provider_for_genai_prices(provider_id: str | None) -> str | None:
    """Map internal provider aliases to genai-prices provider IDs."""
    if provider_id == "bedrock":
        return "aws"
    return provider_id


def get_genai_price_lookup_candidates(
    model_name: str, default_provider: str | None = None
) -> list[tuple[str, str | None]]:
    """
    Build ordered (model_ref, provider_id) candidates for genai-prices lookup.

    This handles common provider/model normalization differences, especially for
    Bedrock Anthropic IDs where runtime model names can be unscoped
    (``anthropic...`` or ``us.anthropic...``) while genai-prices expects
    ``regional.anthropic...`` or ``global.anthropic...``.
    """
    bare_model = strip_provider_prefix(model_name)
    inferred_provider = infer_provider(model_name, default=default_provider)
    provider_id = _normalize_provider_for_genai_prices(inferred_provider)

    candidates: list[tuple[str, str | None]] = [(bare_model, provider_id)]
    if provider_id == "aws":
        unscoped_model = bare_model
        if bare_model.startswith(("us.", "eu.")) and "." in bare_model:
            unscoped_model = bare_model.split(".", 1)[1]
            candidates.append((unscoped_model, provider_id))

        if (
            unscoped_model.startswith("anthropic.")
            and not unscoped_model.startswith("regional.")
            and not unscoped_model.startswith("global.")
        ):
            candidates.append((f"regional.{unscoped_model}", provider_id))
            candidates.append((f"global.{unscoped_model}", provider_id))

    deduped_candidates: list[tuple[str, str | None]] = []
    seen: set[tuple[str, str | None]] = set()
    for candidate in candidates:
        if candidate in seen:
            continue
        deduped_candidates.append(candidate)
        seen.add(candidate)
    return deduped_candidates


def calculate_cost(
    model_name: str,
    input_tokens: int | Any = 0,  # noqa: ANN401
    output_tokens: int | Any = 0,  # noqa: ANN401
    reasoning_tokens: int | Any = 0,  # noqa: ANN401
    cache_read_input_tokens: int | Any = 0,  # noqa: ANN401
    cache_creation_input_tokens: int | Any = 0,  # noqa: ANN401
) -> float:
    """
    Calculate the cost for a model usage using genai-prices with cache-aware pricing.

    genai-prices applies provider-specific discounted/premium rates for cache tokens:
    - Cache reads: ~10% of regular input price (Anthropic), ~50% (OpenAI)
    - Cache creation: ~125% of regular input price (Anthropic)

    IMPORTANT: input_tokens MUST be the TOTAL input token count (including cached tokens).
    Cache tokens are subsets of input_tokens, not additive. For example:
    - OpenAI: prompt_tokens already includes cached_tokens
    - pydantic-ai: input_tokens already includes cache_read + cache_write
    - Anthropic raw API: input_tokens is non-cached only, so callers must pre-sum
      (input_tokens + cache_read + cache_creation) before passing to this function.

    Args:
        model_name: Name of the model (e.g. "gpt-4o", "claude-sonnet-4-20250514",
            or with provider prefix "openai:gpt-4o").
        input_tokens: TOTAL number of input/prompt tokens (including cached tokens).
            Cache tokens are a subset of this value, not additive.
        output_tokens: Number of output tokens
        reasoning_tokens: Number of reasoning tokens (added to output tokens)
        cache_read_input_tokens: Number of cached input tokens read (subset of input_tokens)
        cache_creation_input_tokens: Number of tokens written to cache (subset of input_tokens)

    Returns:
        Total cost in USD
    """
    if not GENAI_PRICES_AVAILABLE:
        return 0.0

    # Safely convert all inputs to integers
    input_tokens = _safe_int(input_tokens)
    output_tokens = _safe_int(output_tokens)
    reasoning_tokens = _safe_int(reasoning_tokens)
    cache_read_input_tokens = _safe_int(cache_read_input_tokens)
    cache_creation_input_tokens = _safe_int(cache_creation_input_tokens)

    # input_tokens is already the TOTAL (including cached tokens).
    # Cache tokens are subsets -- do NOT add them again.
    total_input_tokens = input_tokens
    total_output_tokens = output_tokens + reasoning_tokens

    # Sanity check: cache tokens should be subsets of input_tokens.
    # If this fires, a caller is not normalizing tokens correctly (e.g. passing
    # Anthropic's non-cached input_tokens without pre-summing cache tokens).
    if input_tokens > 0 and cache_read_input_tokens + cache_creation_input_tokens > input_tokens:
        logger.warning(
            "Cache tokens exceed input_tokens (possible normalization bug)",
            extra={
                "model": model_name,
                "input_tokens": input_tokens,
                "cache_read_input_tokens": cache_read_input_tokens,
                "cache_creation_input_tokens": cache_creation_input_tokens,
            },
        )

    try:
        usage = GenAIUsage(
            input_tokens=total_input_tokens,
            output_tokens=total_output_tokens,
            cache_read_tokens=cache_read_input_tokens,
            cache_write_tokens=cache_creation_input_tokens,
        )
        last_lookup_error: LookupError | None = None
        for model_ref, provider_id in get_genai_price_lookup_candidates(model_name):
            try:
                price = genai_calc_price(usage, model_ref=model_ref, provider_id=provider_id)
                return float(round(price.total_price, 6))
            except LookupError as lookup_error:
                last_lookup_error = lookup_error
                continue
        if last_lookup_error is not None:
            raise last_lookup_error
        return 0.0
    except Exception as e:
        logger.warning("Failed to calculate cost for %s: %s", model_name, e)
        return 0.0
