"""
LLM Pricing Data Fetcher

Fetches real-time pricing data from various LLM providers.
Supports multiple providers and fallback to static pricing.
"""

from typing import Dict, Optional

import aiohttp
from loguru import logger

from ...llm.providers.base import normalize_provider


class PricingFetcher:
    """Fetches LLM pricing data from various sources."""

    def __init__(self):
        self.session = None
        self.static_pricing = self._load_static_pricing()

    async def __aenter__(self):
        self.session = aiohttp.ClientSession()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()

    def _load_static_pricing(self) -> Dict[str, Dict[str, float]]:
        """Load static pricing as fallback (per 1K tokens in USD)."""
        return {
            # OpenAI Models
            "gpt-4": {"prompt": 0.03, "completion": 0.06},
            "gpt-4-turbo": {"prompt": 0.01, "completion": 0.03},
            "gpt-4o": {"prompt": 0.005, "completion": 0.015},
            "gpt-4o-mini": {"prompt": 0.00015, "completion": 0.0006},
            "gpt-3.5-turbo": {"prompt": 0.001, "completion": 0.002},
            # Anthropic Models
            "claude-3-opus": {"prompt": 0.015, "completion": 0.075},
            "claude-3-sonnet": {"prompt": 0.003, "completion": 0.015},
            "claude-3-haiku": {"prompt": 0.00025, "completion": 0.00125},
            "claude-3-5-sonnet": {"prompt": 0.003, "completion": 0.015},
            # Google Models
            "gemini-pro": {"prompt": 0.0005, "completion": 0.0015},
            "gemini-pro-vision": {"prompt": 0.0005, "completion": 0.0015},
            "gemini-1.5-pro": {"prompt": 0.0035, "completion": 0.0105},
            "gemini-1.5-flash": {"prompt": 0.00035, "completion": 0.00105},
            # Local/Open Source (free)
            "ollama": {"prompt": 0.0, "completion": 0.0},
            "llama": {"prompt": 0.0, "completion": 0.0},
            "mistral": {"prompt": 0.0, "completion": 0.0},
            "gemma": {"prompt": 0.0, "completion": 0.0},
            "qwen": {"prompt": 0.0, "completion": 0.0},
            "codellama": {"prompt": 0.0, "completion": 0.0},
            "vicuna": {"prompt": 0.0, "completion": 0.0},
            "alpaca": {"prompt": 0.0, "completion": 0.0},
            "lmstudio": {"prompt": 0.0, "completion": 0.0},
            "llamacpp": {"prompt": 0.0, "completion": 0.0},
        }

    async def get_model_pricing(
        self, model_name: str, provider: str = None
    ) -> Optional[Dict[str, float]]:
        """Get pricing for a specific model and provider."""
        # Normalize inputs
        model_name = model_name.lower() if model_name else ""
        provider = normalize_provider(provider) or ""

        # Provider-first approach: Check if provider indicates local/free models
        local_providers = ["ollama", "lmstudio", "llamacpp"]
        if provider in local_providers:
            logger.debug(
                f"Local provider '{provider}' detected - returning zero cost"
            )
            return {"prompt": 0.0, "completion": 0.0}

        # Fallback to static pricing with provider priority
        if provider:
            # First try provider-specific lookup with exact matching
            provider_models = self._get_models_by_provider(provider)
            # Try exact match
            if model_name in provider_models:
                return provider_models[model_name]
            # Try exact match without provider prefix
            if "/" in model_name:
                model_only = model_name.split("/")[-1]
                if model_only in provider_models:
                    return provider_models[model_only]

        # Exact model name matching only
        # First try exact match
        if model_name in self.static_pricing:
            return self.static_pricing[model_name]

        # Try exact match without provider prefix (e.g., "openai/gpt-4o-mini" -> "gpt-4o-mini")
        if "/" in model_name:
            model_only = model_name.split("/")[-1]
            if model_only in self.static_pricing:
                return self.static_pricing[model_only]

        # No pricing found - return None instead of default pricing
        logger.warning(
            f"No pricing found for model: {model_name}, provider: {provider}"
        )
        return None

    def _get_models_by_provider(
        self, provider: str
    ) -> Dict[str, Dict[str, float]]:
        """Get models for a specific provider."""
        provider = normalize_provider(provider) or ""
        provider_models = {}

        if provider == "openai":
            provider_models = {
                k: v
                for k, v in self.static_pricing.items()
                if k.startswith("gpt")
            }
        elif provider == "anthropic":
            provider_models = {
                k: v
                for k, v in self.static_pricing.items()
                if k.startswith("claude")
            }
        elif provider == "google":
            provider_models = {
                k: v
                for k, v in self.static_pricing.items()
                if k.startswith("gemini")
            }
        elif provider in ["ollama", "lmstudio", "llamacpp"]:
            # All local models are free
            provider_models = {
                k: v
                for k, v in self.static_pricing.items()
                if v["prompt"] == 0.0 and v["completion"] == 0.0
            }

        return provider_models

    def get_provider_from_model(self, model_name: str) -> str:
        """Determine the provider from model name."""
        model_name = model_name.lower()

        if "gpt" in model_name or "openai" in model_name:
            return "openai"
        if "claude" in model_name or "anthropic" in model_name:
            return "anthropic"
        if "gemini" in model_name or "google" in model_name:
            return "google"
        if "llama" in model_name or "meta" in model_name:
            return "meta"
        if "mistral" in model_name:
            return "mistral"
        if "ollama" in model_name:
            return "ollama"
        return "unknown"
