"""
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

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 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:
            logger.warning("Failed to fetch OpenAI pricing")
            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:
            logger.warning("Failed to fetch Anthropic pricing")
            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:
            logger.warning("Failed to fetch Google pricing")
            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:
            logger.warning("Failed to fetch HuggingFace pricing")
            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 = 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}

        # 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 = 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

    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"
        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"
