"""Pricing lookup for deployed models using a 3-tier strategy."""

from __future__ import annotations

import logging

import litellm

from sage_evaluator.models import ModelPricing, PricingSource

logger = logging.getLogger(__name__)

# Common prefixes that should be stripped before doing a bare-name lookup.
_STRIP_PREFIXES: tuple[str, ...] = ("azure_ai/", "azure/", "openai/")

# Per-token costs for common models when litellm has no entry.
# Tuples are (input_cost_per_token, output_cost_per_token).
FALLBACK_PRICING: dict[str, tuple[float, float]] = {
    "gpt-4o": (0.0000025, 0.00001),
    "gpt-4o-mini": (0.00000015, 0.0000006),
    "gpt-4": (0.00003, 0.00006),
    "gpt-3.5-turbo": (0.0000005, 0.0000015),
    "claude-opus-4-6": (0.000015, 0.000075),
    "claude-sonnet-4-20250514": (0.000003, 0.000015),
    "claude-3-5-haiku-20241022": (0.0000008, 0.000004),
}


class PricingLookup:
    """Resolve per-token pricing for a model using a 3-tier fallback strategy.

    Tier 1 — litellm's ``model_cost`` dictionary (exact key or with common
              provider prefixes like ``azure_ai/`` and ``azure/`` prepended).
    Tier 2 — ``FALLBACK_PRICING`` hard-coded table for common Azure AI models.
    Tier 3 — Zero costs with ``PricingSource.UNKNOWN``.

    The lookup is intentionally non-raising: callers receive a valid
    ``ModelPricing`` object regardless of whether pricing data is available.

    Example::

        lookup = PricingLookup()
        pricing = lookup.get_pricing("azure_ai/gpt-4o")
        # ModelPricing(model='azure_ai/gpt-4o', input_cost_per_token=2.5e-06, ...)
    """

    def get_pricing(self, model: str) -> ModelPricing:
        """Return pricing for *model*.  Never raises.

        Args:
            model: Model identifier, e.g. ``"azure_ai/gpt-4o"`` or ``"gpt-4o"``.

        Returns:
            A :class:`~sage_evaluator.models.ModelPricing` instance whose
            ``source`` field indicates which pricing tier was used.
        """
        # --- Tier 1: litellm exact match (try original key, then with azure/
        #             prefix, then with azure_ai/ prefix) ---
        candidates = self._build_litellm_candidates(model)
        for candidate in candidates:
            entry = litellm.model_cost.get(candidate)
            if entry and "input_cost_per_token" in entry and "output_cost_per_token" in entry:
                logger.debug("Pricing for %r sourced from litellm (key=%r)", model, candidate)
                return ModelPricing(
                    model=model,
                    input_cost_per_token=float(entry["input_cost_per_token"]),
                    output_cost_per_token=float(entry["output_cost_per_token"]),
                    source=PricingSource.LITELLM,
                )

        # --- Tier 2: hard-coded fallback table ---
        bare_name = self._strip_prefix(model)
        if bare_name in FALLBACK_PRICING:
            input_cost, output_cost = FALLBACK_PRICING[bare_name]
            logger.debug("Pricing for %r sourced from fallback table", model)
            return ModelPricing(
                model=model,
                input_cost_per_token=input_cost,
                output_cost_per_token=output_cost,
                source=PricingSource.FALLBACK,
            )

        # --- Tier 3: unknown ---
        logger.debug("No pricing found for %r; returning zeros", model)
        return ModelPricing(
            model=model,
            input_cost_per_token=0.0,
            output_cost_per_token=0.0,
            source=PricingSource.UNKNOWN,
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _strip_prefix(model: str) -> str:
        """Return *model* with the first matching provider prefix removed."""
        for prefix in _STRIP_PREFIXES:
            if model.startswith(prefix):
                return model[len(prefix) :]
        return model

    @staticmethod
    def _build_litellm_candidates(model: str) -> list[str]:
        """Return an ordered list of keys to probe in ``litellm.model_cost``.

        The list starts with the original model string (handles already-prefixed
        keys like ``azure_ai/gpt-4o`` as well as plain ``gpt-4o``), then tries
        the bare name with ``azure/`` prepended (common for Azure OpenAI), and
        finally the bare name without any prefix.
        """
        bare = PricingLookup._strip_prefix(model)
        candidates: list[str] = [model]
        if bare != model:
            # e.g. azure_ai/gpt-4o -> also try azure/gpt-4o and gpt-4o
            for prefix in _STRIP_PREFIXES:
                prefixed = f"{prefix}{bare}"
                if prefixed != model:
                    candidates.append(prefixed)
            candidates.append(bare)
        else:
            # Original was already bare; try common prefixes too
            for prefix in _STRIP_PREFIXES:
                candidates.append(f"{prefix}{bare}")
        return candidates
