"""
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
import threading
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.")

try:
    from genai_prices.update_prices import UpdatePrices
except ImportError:
    UpdatePrices = None

_price_updater: Any = None
_price_updater_lock = threading.Lock()


def enable_live_price_updates(*, update_interval: float = 3600) -> None:
    """Opt in to background price updates from the genai-prices GitHub repo.

    Call this once at application startup (e.g. after ``AmplitudeAI()`` init)
    to fetch the latest pricing data every *update_interval* seconds.  This
    ensures new model pricing is available within days of being added to the
    genai-prices repository, instead of waiting for a PyPI package release.

    This makes outbound HTTPS requests to raw.githubusercontent.com.
    Only enable in environments where outbound network access is permitted.
    """
    global _price_updater  # noqa: PLW0603
    if _price_updater is not None or not GENAI_PRICES_AVAILABLE or UpdatePrices is None:
        return
    with _price_updater_lock:
        if _price_updater is not None:
            return
        try:
            _price_updater = UpdatePrices(update_interval=update_interval)
            _price_updater.start(wait=False)
        except Exception:
            logger.debug("Failed to start live price updater", exc_info=True)


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 _has_provider_prefix(model_name: str) -> bool:
    """Check if model_name has a provider prefix (e.g. 'openai:gpt-4o').

    Provider prefixes are simple identifiers without dots. Bedrock model IDs
    like 'anthropic.claude-sonnet-4-20250514-v1:0' contain ':' in the version
    suffix but NOT as a provider separator — the part before ':' contains dots.
    """
    if ":" not in model_name:
        return False
    prefix = model_name.split(":", 1)[0]
    return "." not in prefix


def strip_provider_prefix(model_name: str) -> str:
    """Strip provider prefix from model name (e.g. 'openai:gpt-4o' -> 'gpt-4o')."""
    if _has_provider_prefix(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 _has_provider_prefix(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"
    if provider_id == "gemini":
        return "google"
    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.

    For AWS/Bedrock models, uses a **generalized** dot-prefix stripping strategy
    instead of enumerating known regions or vendors.  Bedrock model IDs follow
    ``[region.][vendor.]model-name[-version]`` — we progressively strip
    dot-separated prefixes and try each variant with and without provider,
    plus ``regional.`` / ``global.`` prefixes that genai-prices uses.

    This approach is forward-compatible: new AWS regions and Bedrock vendors
    work automatically without code changes.
    """
    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)

    is_bedrock = provider_id == "aws"

    candidates: list[tuple[str, str | None]] = [(bare_model, provider_id)]
    # For Bedrock, also try without provider for globally-matched models (e.g. Claude)
    if is_bedrock:
        candidates.append((bare_model, None))

    # For any model with dot-separated segments (e.g. vendor.model, region.vendor.model),
    # progressively strip prefixes and try each variant.  Safe because calculate_cost
    # iterates candidates and stops at the first successful price hit.
    # For Bedrock specifically, also try regional./global. prefixes.
    if "." in bare_model:
        parts = bare_model.split(".")
        for i in range(1, len(parts)):
            stripped_part = ".".join(parts[i:])
            candidates.append((stripped_part, provider_id))
            candidates.append((stripped_part, None))

        if is_bedrock:
            # genai-prices often indexes Bedrock models under regional.X / global.X
            vendor_model = bare_model
            first_seg = parts[0]
            if first_seg not in ("regional", "global") and len(parts) > 2:
                vendor_model = ".".join(parts[1:])
            if not vendor_model.startswith(("regional.", "global.")):
                candidates.append((f"regional.{vendor_model}", provider_id))
                candidates.append((f"global.{vendor_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
    default_provider: str | None = None,
) -> 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: IGNORED — kept for backward compatibility. Reasoning
            tokens are already included in output_tokens by providers (e.g.
            OpenAI's completion_tokens includes reasoning_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)
        default_provider: Fallback provider when model_name has no prefix
            (e.g. "bedrock" for OTEL spans where provider is a separate attribute).

    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.
    # output_tokens already includes reasoning tokens (e.g. OpenAI's
    # completion_tokens includes reasoning_tokens) -- do NOT add them again.
    total_input_tokens = input_tokens
    total_output_tokens = output_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, default_provider=default_provider):
            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
