"""
Provider Fallback System for AIIQ
Cycles through multiple LLM providers until one succeeds
"""
import asyncio
import logging
from dataclasses import dataclass
from typing import Any, Optional, Dict, List
from enum import Enum
import os

try:
    import litellm
except ImportError:
    litellm = None

logger = logging.getLogger(__name__)


class ProviderStatus(Enum):
    """Provider connection status"""
    SUCCESS = "success"
    INVALID_KEY = "invalid_key"
    QUOTA_EXCEEDED = "quota_exceeded"
    NOT_FOUND = "not_found"
    CONNECTION_ERROR = "connection_error"
    RATE_LIMIT = "rate_limit"
    UNKNOWN_ERROR = "unknown_error"


@dataclass
class ProviderConfig:
    """Configuration for a single provider"""
    name: str
    provider: str
    api_key: str
    api_base: Optional[str] = None
    model: Optional[str] = None
    enabled: bool = True
    priority: int = 100


@dataclass
class ProviderAttempt:
    """Result of a provider attempt"""
    provider_name: str
    status: ProviderStatus
    error_message: Optional[str] = None
    response: Optional[Any] = None
    latency_seconds: float = 0.0


class ProviderFallbackService:
    """Manages provider fallback logic with automatic cycling"""

    def __init__(self):
        self.providers: List[ProviderConfig] = []
        self._load_providers_from_env()
        self._sort_providers()

    def _get_env_any(self, *names: str) -> Optional[str]:
        """Get environment value from any of the provided names"""
        for name in names:
            value = os.getenv(name)
            if value:
                return value
        return None

    def _load_providers_from_env(self):
        """Load provider configurations from environment variables"""

        # OpenRouter - check both BASE_URL and API_BASE variants
        openrouter_key = os.getenv("OPENROUTER_API_KEY")
        if openrouter_key:
            openrouter_base = self._get_env_any("OPENROUTER_API_BASE", "OPENROUTER_BASE_URL")
            openrouter_model = os.getenv("OPENROUTER_MODEL", "meta-llama/llama-3.3-70b-instruct:free")
            
            self.providers.append(ProviderConfig(
                name="openrouter",
                provider="openrouter",
                api_key=openrouter_key,
                api_base=openrouter_base or "https://openrouter.ai/api/v1",
                model=openrouter_model,
                priority=10
            ))

        # ZAI - check both ENDPOINT and API_BASE variants
        zai_key = self._get_env_any("ZAI_API_KEY", "ZAI_KEY")
        if zai_key:
            zai_base = self._get_env_any("ZAI_API_BASE", "ZAI_ENDPOINT", "ZAI_BASE_URL")
            zai_model = os.getenv("ZAI_MODEL", "glm-4.7")
            
            self.providers.append(ProviderConfig(
                name="zai",
                provider="z_ai",  # litellm uses "z_ai" not "zai"
                api_key=zai_key,
                api_base=zai_base or "https://api.z.ai/api/coding/paas/v4",
                model=zai_model,
                priority=20
            ))

        # Ollama
        ollama_base = self._get_env_any("OLLAMA_API_BASE", "OLLAMA_BASE_URL")
        if ollama_base:
            self.providers.append(ProviderConfig(
                name="ollama",
                provider="ollama",
                api_key="ollama",  # Ollama doesn't need a real key
                api_base=ollama_base,
                model=os.getenv("OLLAMA_MODEL", "llama3.2"),
                priority=30
            ))

        # OpenAI
        openai_key = os.getenv("OPENAI_API_KEY")
        if openai_key:
            openai_base = self._get_env_any("OPENAI_API_BASE", "OPENAI_BASE_URL")
            
            self.providers.append(ProviderConfig(
                name="openai",
                provider="openai",
                api_key=openai_key,
                api_base=openai_base or "https://api.openai.com/v1",
                model=os.getenv("OPENAI_MODEL", "gpt-4o-mini"),
                priority=40
            ))

        # Anthropic
        anthropic_key = os.getenv("ANTHROPIC_API_KEY")
        if anthropic_key:
            anthropic_base = self._get_env_any("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL")
            
            self.providers.append(ProviderConfig(
                name="anthropic",
                provider="anthropic",
                api_key=anthropic_key,
                api_base=anthropic_base or "https://api.anthropic.com",
                model=os.getenv("ANTHROPIC_MODEL", "claude-3-haiku-20240307"),
                priority=50
            ))

        # Gemini
        gemini_key = os.getenv("GEMINI_API_KEY")
        if gemini_key:
            gemini_base = self._get_env_any("GEMINI_API_BASE", "GEMINI_BASE_URL")
            
            self.providers.append(ProviderConfig(
                name="gemini",
                provider="gemini",
                api_key=gemini_key,
                api_base=gemini_base or "https://generativelanguage.googleapis.com",
                model=os.getenv("GEMINI_MODEL", "gemini/gemini-2.0-flash-exp"),
                priority=60
            ))

    def _sort_providers(self):
        """Sort providers by priority (lower first)"""
        self.providers.sort(key=lambda p: p.priority)

    def _classify_error(self, error: Exception) -> ProviderStatus:
        """Classify an error to determine if we should retry"""
        error_str = str(error).lower()

        if any(term in error_str for term in ["invalid", "unauthorized", "401", "authentication"]):
            return ProviderStatus.INVALID_KEY
        if any(term in error_str for term in ["quota", "limit exceeded", "429", "insufficient"]):
            return ProviderStatus.QUOTA_EXCEEDED
        if any(term in error_str for term in ["not found", "404", "does not exist"]):
            return ProviderStatus.NOT_FOUND
        if any(term in error_str for term in ["connection", "timeout", "refused", "network"]):
            return ProviderStatus.CONNECTION_ERROR
        if "rate limit" in error_str:
            return ProviderStatus.RATE_LIMIT

        return ProviderStatus.UNKNOWN_ERROR

    def _should_try_next_provider(self, status: ProviderStatus) -> bool:
        """Determine if we should try the next provider based on error type"""
        return status in [
            ProviderStatus.INVALID_KEY,
            ProviderStatus.QUOTA_EXCEEDED,
            ProviderStatus.NOT_FOUND,
            ProviderStatus.CONNECTION_ERROR,
            ProviderStatus.RATE_LIMIT,
            ProviderStatus.UNKNOWN_ERROR
        ]

    def _format_model_for_litellm(self, provider_config: ProviderConfig) -> str:
        """Format model name for litellm (add provider prefix if needed)"""
        model = provider_config.model
        
        # For some providers, litellm needs the provider/ prefix
        if provider_config.provider == "openrouter":
            # OpenRouter models should be prefixed
            if not model.startswith("openrouter/"):
                return f"openrouter/{model}"
        elif provider_config.provider == "z_ai":
            # Z.AI might need prefix
            if not model.startswith("z_ai/"):
                return f"z_ai/{model}"
                
        return model

    async def complete_with_fallback(
        self,
        messages: list,
        fallback_model: Optional[str] = None,
        **kwargs
    ) -> tuple[Optional[Any], list[ProviderAttempt]]:
        """Try completing with messages, cycling through providers until success."""
        if not litellm:
            logger.error("litellm not installed - cannot use provider fallback")
            return None, []

        if not self.providers:
            logger.warning("No providers configured")
            return None, []

        attempts: list[ProviderAttempt] = []

        for provider_config in self.providers:
            if not provider_config.enabled:
                continue

            logger.info(f"Trying provider: {provider_config.name}")

            try:
                import time
                start_time = time.time()

                # Format model for litellm
                model = self._format_model_for_litellm(provider_config)
                
                litellm_kwargs = {
                    "model": model,
                    "messages": messages,
                    "api_key": provider_config.api_key,
                }

                if provider_config.api_base:
                    litellm_kwargs["api_base"] = provider_config.api_base
                
                # Add any additional kwargs
                litellm_kwargs.update(kwargs)

                logger.debug(f"Calling litellm with model={model}, provider={provider_config.provider}")

                response = await asyncio.to_thread(
                    litellm.completion,
                    **litellm_kwargs
                )

                latency = time.time() - start_time

                attempt = ProviderAttempt(
                    provider_name=provider_config.name,
                    status=ProviderStatus.SUCCESS,
                    response=response,
                    latency_seconds=latency
                )
                attempts.append(attempt)

                logger.info(f"[OK] Provider {provider_config.name} succeeded in {latency:.2f}s")
                return response, attempts

            except Exception as e:
                latency = time.time() - start_time
                status = self._classify_error(e)

                attempt = ProviderAttempt(
                    provider_name=provider_config.name,
                    status=status,
                    error_message=str(e)[:500],
                    latency_seconds=latency
                )
                attempts.append(attempt)

                logger.warning(f"[FAIL] Provider {provider_config.name} failed: {status.value} - {str(e)[:200]}")

                if not self._should_try_next_provider(status):
                    break

        logger.error(f"All {len(attempts)} providers failed")
        return None, attempts

    def list_providers(self) -> list[Dict[str, Any]]:
        """List all configured providers with their status"""
        return [
            {
                "name": p.name,
                "provider": p.provider,
                "model": p.model,
                "api_base": p.api_base,
                "enabled": p.enabled,
                "priority": p.priority,
                "has_key": bool(p.api_key and p.api_key != "" and p.api_key != "ollama")
            }
            for p in self.providers
        ]


_fallback_service: Optional[ProviderFallbackService] = None


def get_provider_fallback_service() -> ProviderFallbackService:
    """Get the singleton provider fallback service"""
    global _fallback_service
    if _fallback_service is None:
        _fallback_service = ProviderFallbackService()
    return _fallback_service


async def complete_with_fallback(messages: list, **kwargs) -> Optional[Any]:
    """Convenience function: complete messages with automatic provider fallback."""
    service = get_provider_fallback_service()
    response, attempts = await service.complete_with_fallback(messages, **kwargs)

    if response is None:
        logger.error(
            f"All {len(attempts)} provider attempts failed:\n" +
            "\n".join([
                f"  - {a.provider_name}: {a.status.value}"
                for a in attempts
            ])
        )

    return response
