"""
LLM Pricing Data Fetcher

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

from typing import Any, Dict, Optional

import aiohttp
from loguru import logger


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},
            "vllm": {"prompt": 0.0, "completion": 0.0},
            "lmstudio": {"prompt": 0.0, "completion": 0.0},
            "llamacpp": {"prompt": 0.0, "completion": 0.0},
        }

    async def fetch_openai_pricing(self) -> Optional[Dict[str, Any]]:
        """Fetch OpenAI pricing from their API (if available)."""
        try:
            # Note: OpenAI doesn't have a public pricing API
            # This would need to be web scraping or manual updates
            logger.info("Using static OpenAI pricing (no public API available)")
            return None
        except Exception as e:
            logger.warning(f"Failed to fetch OpenAI pricing: {e}")
            return None

    async def fetch_anthropic_pricing(self) -> Optional[Dict[str, Any]]:
        """Fetch Anthropic pricing."""
        try:
            # Note: Anthropic doesn't have a public pricing API
            # This would need to be web scraping or manual updates
            logger.info(
                "Using static Anthropic pricing (no public API available)"
            )
            return None
        except Exception as e:
            logger.warning(f"Failed to fetch Anthropic pricing: {e}")
            return None

    async def fetch_google_pricing(self) -> Optional[Dict[str, Any]]:
        """Fetch Google/Gemini pricing."""
        try:
            # Note: Google doesn't have a dedicated pricing API for individual models
            # This would need to be web scraping or manual updates
            logger.info("Using static Google pricing (no public API available)")
            return None
        except Exception as e:
            logger.warning(f"Failed to fetch Google pricing: {e}")
            return None

    async def fetch_huggingface_pricing(self) -> Optional[Dict[str, Any]]:
        """Fetch HuggingFace Inference API pricing."""
        try:
            if not self.session:
                return None

            # HuggingFace has some pricing info but not a structured API
            # This is more for hosted inference endpoints
            url = "https://huggingface.co/pricing"
            async with self.session.get(url) as response:
                if response.status == 200:
                    # Would need to parse HTML for pricing info
                    logger.info(
                        "HuggingFace pricing would require HTML parsing"
                    )
                    return None
        except Exception as e:
            logger.warning(f"Failed to fetch HuggingFace pricing: {e}")
            return None

    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 = provider.lower() if provider else ""

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

        # Try to fetch live pricing first (most providers don't have APIs)
        if (
            provider == "openai"
            or "gpt" in model_name
            or "openai" in model_name
        ):
            await self.fetch_openai_pricing()
        elif (
            provider == "anthropic"
            or "claude" in model_name
            or "anthropic" in model_name
        ):
            await self.fetch_anthropic_pricing()
        elif (
            provider == "google"
            or "gemini" in model_name
            or "google" in model_name
        ):
            await self.fetch_google_pricing()

        # 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 = provider.lower()
        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", "vllm", "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

    async def get_all_pricing(self) -> Dict[str, Dict[str, float]]:
        """Get pricing for all known models."""
        # In the future, this could aggregate from multiple live sources
        return self.static_pricing.copy()

    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"
        elif "claude" in model_name or "anthropic" in model_name:
            return "anthropic"
        elif "gemini" in model_name or "google" in model_name:
            return "google"
        elif "llama" in model_name or "meta" in model_name:
            return "meta"
        elif "mistral" in model_name:
            return "mistral"
        elif "ollama" in model_name:
            return "ollama"
        else:
            return "unknown"
