"""Unified pricing registry and resolver for all LLM and embedding models.

This module centralizes model pricing data for OpenAI and Azure OpenAI
clients (both chat completion and embedding), enabling single-source-of-truth
pricing maintenance and consistent fallback behavior across all clients.

Shared pricing constants reduce duplication for models where OpenAI and Azure
pricing is identical. Provider-specific registries maintain semantic clarity
about where each price applies.

Usage:
    # For OpenAI chat completion clients
    pricing = PricingRegistry.resolve_openai_chat_pricing("gpt-4.1-mini")

    # For Azure OpenAI chat completion clients
    pricing = PricingRegistry.resolve_azure_chat_pricing("gpt-4.1-mini")

    # For OpenAI embedding clients
    pricing = PricingRegistry.resolve_openai_embedding_pricing("text-embedding-3-large")

    # For Azure OpenAI embedding clients
    pricing = PricingRegistry.resolve_azure_embedding_pricing("text-embedding-3-large")
"""

import logging

from .types import ModelPricing

logger = logging.getLogger(__name__)


# ===== Shared Pricing Constants (same across OpenAI and Azure) =====
# Chat models with identical pricing on both platforms
_SHARED_GPT4_PRICING = ModelPricing(input_per_1m=30.0, output_per_1m=60.0)  # $0.03/$0.06 per 1K
_SHARED_GPT4_TURBO_PRICING = ModelPricing(input_per_1m=10.0, output_per_1m=30.0)  # $0.01/$0.03 per 1K
_SHARED_GPT35_PRICING = ModelPricing(input_per_1m=0.50, output_per_1m=1.50)
_SHARED_GPT4O_MINI_PRICING = ModelPricing(input_per_1m=0.15, output_per_1m=0.60)

# Embedding models with identical pricing on both platforms
_SHARED_EMBEDDING_3_LARGE = ModelPricing(input_per_1m=0.13, output_per_1m=0.0)
_SHARED_EMBEDDING_3_SMALL = ModelPricing(input_per_1m=0.02, output_per_1m=0.0)
_SHARED_EMBEDDING_ADA_002 = ModelPricing(input_per_1m=0.10, output_per_1m=0.0)


class PricingRegistry:
    """Central registry of model pricing with resolver functions for all providers."""

    # ===== OpenAI Chat Model Pricing (per 1M tokens) =====
    OPENAI_CHAT_PRICING: dict[str, ModelPricing] = {
        "gpt-4.1-mini": ModelPricing(input_per_1m=0.15, output_per_1m=0.60),
        "gpt-4.1-nano": ModelPricing(input_per_1m=0.10, output_per_1m=0.40),
        "gpt-4.1": ModelPricing(input_per_1m=0.40, output_per_1m=1.60),
        "gpt-4o-mini": _SHARED_GPT4O_MINI_PRICING,
        "gpt-4": _SHARED_GPT4_PRICING,
        "gpt-4-turbo": _SHARED_GPT4_TURBO_PRICING,
        "gpt-3.5-turbo": _SHARED_GPT35_PRICING,
        "gpt-5": ModelPricing(input_per_1m=2.50, output_per_1m=20.00),
        "gpt-5-mini": ModelPricing(input_per_1m=0.25, output_per_1m=2.00),
        "gpt-5.1": ModelPricing(input_per_1m=1.25, output_per_1m=10.00),
        "gpt-5.2": ModelPricing(input_per_1m=1.75, output_per_1m=7.00),
        "gpt-5.2-mini": ModelPricing(input_per_1m=0.45, output_per_1m=3.60),
    }

    # ===== Azure OpenAI Chat Model Pricing (per 1M tokens) =====
    AZURE_OPENAI_CHAT_PRICING: dict[str, ModelPricing] = {
        "gpt-4.1-mini": ModelPricing(input_per_1m=0.40, output_per_1m=1.60),
        "gpt-4o-mini": _SHARED_GPT4O_MINI_PRICING,
        "gpt-4": _SHARED_GPT4_PRICING,
        "gpt-4-turbo": _SHARED_GPT4_TURBO_PRICING,
        "gpt-3.5-turbo": _SHARED_GPT35_PRICING,
        "gpt-4.1": ModelPricing(input_per_1m=2.0, output_per_1m=8.0),
        "gpt-4.1-nano": ModelPricing(input_per_1m=0.10, output_per_1m=0.40),
        "gpt-5-mini": ModelPricing(input_per_1m=0.25, output_per_1m=2.0),
        "gpt-5.1-chat": ModelPricing(input_per_1m=1.25, output_per_1m=10.0),
    }

    # ===== OpenAI Embedding Pricing (per 1M tokens input, output always 0) =====
    OPENAI_EMBEDDING_PRICING: dict[str, ModelPricing] = {
        "text-embedding-3-large": _SHARED_EMBEDDING_3_LARGE,
        "text-embedding-3-small": _SHARED_EMBEDDING_3_SMALL,
        "text-embedding-ada-002": _SHARED_EMBEDDING_ADA_002,
    }

    # ===== Azure OpenAI Embedding Pricing (per 1M tokens input, output always 0) =====
    AZURE_OPENAI_EMBEDDING_PRICING: dict[str, ModelPricing] = {
        "text-embedding-3-large": _SHARED_EMBEDDING_3_LARGE,
        "text-embedding-3-small": _SHARED_EMBEDDING_3_SMALL,
        "text-embedding-ada-002": _SHARED_EMBEDDING_ADA_002,
    }

    @staticmethod
    def resolve_openai_chat_pricing(
        model: str,
        default_model: str = "gpt-4",
    ) -> ModelPricing:
        """
        Resolve OpenAI chat model pricing by exact match, then prefix match.

        Args:
            model: Model name to resolve (e.g., "gpt-4.1-mini")
            default_model: Fallback model if no match found

        Returns:
            ModelPricing for the model, or default pricing with warning
        """
        return PricingRegistry._resolve(
            model,
            PricingRegistry.OPENAI_CHAT_PRICING,
            default_model,
        )

    @staticmethod
    def resolve_azure_chat_pricing(
        model: str,
        default_model: str = "gpt-4",
    ) -> ModelPricing:
        """
        Resolve Azure OpenAI chat model pricing by exact match, then prefix match.

        Args:
            model: Model name to resolve (e.g., "gpt-4.1-mini")
            default_model: Fallback model if no match found

        Returns:
            ModelPricing for the model, or default pricing with warning
        """
        return PricingRegistry._resolve(
            model,
            PricingRegistry.AZURE_OPENAI_CHAT_PRICING,
            default_model,
        )

    @staticmethod
    def resolve_openai_embedding_pricing(
        model: str,
        default_model: str = "text-embedding-ada-002",
    ) -> ModelPricing:
        """
        Resolve OpenAI embedding model pricing by exact match, then prefix match.

        Args:
            model: Model name to resolve (e.g., "text-embedding-3-large")
            default_model: Fallback model if no match found

        Returns:
            ModelPricing for the model, or default pricing with warning
        """
        return PricingRegistry._resolve(
            model,
            PricingRegistry.OPENAI_EMBEDDING_PRICING,
            default_model,
        )

    @staticmethod
    def resolve_azure_embedding_pricing(
        model: str,
        default_model: str = "text-embedding-ada-002",
    ) -> ModelPricing:
        """
        Resolve Azure OpenAI embedding model pricing by exact match, then prefix match.

        Args:
            model: Model name to resolve (e.g., "text-embedding-3-large")
            default_model: Fallback model if no match found

        Returns:
            ModelPricing for the model, or default pricing with warning
        """
        return PricingRegistry._resolve(
            model,
            PricingRegistry.AZURE_OPENAI_EMBEDDING_PRICING,
            default_model,
        )

    @staticmethod
    def _resolve(
        model: str,
        pricing_map: dict[str, ModelPricing],
        default_model: str,
    ) -> ModelPricing:
        """
        Internal resolver: find model by exact match first, then longest prefix match.

        Strategy:
        1. Try exact match (e.g., "gpt-4.1-mini" matches exactly)
        2. Fall back to longest prefix match (e.g., "gpt-4.1-mini-custom" matches "gpt-4.1-mini")
        3. If neither works, warn and use default pricing

        Args:
            model: Model name to find
            pricing_map: Registry to search (e.g., OPENAI_CHAT_PRICING)
            default_model: Fallback model name in case of no match

        Returns:
            ModelPricing object; warns and falls back to default if model not found
        """
        # Strategy 1: Try exact match first
        if model in pricing_map:
            return pricing_map[model]

        # Strategy 2: Try longest prefix match (e.g., "gpt-4.1-mini" before "gpt-4.1" before "gpt-4")
        for model_key in sorted(pricing_map.keys(), key=len, reverse=True):
            if model.startswith(model_key):
                return pricing_map[model_key]

        # No match: warn and use default
        logger.warning(
            "Model pricing not available for model '%s'. "
            "Using default '%s' pricing. "
            "To override, import ModelPricing and pass custom pricing:\n"
            "  from agentbyte.llm import ModelPricing\n"
            "  custom_pricing = ModelPricing(input_per_1m=0.10, output_per_1m=0.50)",
            model,
            default_model,
        )
        return pricing_map[default_model]


__all__ = ["PricingRegistry"]
