"""
AI Provider Abstraction Layer
Supports: OpenAI, Anthropic, Ollama (Local), with automatic fallback
"""
from typing import Optional, List, Dict, Any, AsyncGenerator
from abc import ABC, abstractmethod
import httpx
import json
from enum import Enum
import logging

logger = logging.getLogger(__name__)


class AIProvider(str, Enum):
    """Available AI providers"""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    OLLAMA = "ollama"  # Local AI, no API key needed
    CUSTOM = "custom"


class AIMessage:
    """Standardized AI message format"""
    def __init__(self, role: str, content: str):
        self.role = role
        self.content = content

    def to_dict(self) -> Dict[str, str]:
        return {"role": self.role, "content": self.content}


class BaseAIProvider(ABC):
    """Base class for all AI providers"""

    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None):
        self.api_key = api_key
        self.base_url = base_url
        self.available = False

    @abstractmethod
    async def chat(
        self,
        messages: List[AIMessage],
        model: str,
        temperature: float = 0.7,
        max_tokens: int = 1000,
        stream: bool = False
    ) -> str | AsyncGenerator[str, None]:
        """Send chat completion request"""
        pass

    @abstractmethod
    async def is_available(self) -> bool:
        """Check if provider is available and configured"""
        pass


class OpenAIProvider(BaseAIProvider):
    """OpenAI API provider (GPT-4, GPT-3.5, etc.)"""

    def __init__(self, api_key: str, base_url: str = "https://api.openai.com/v1"):
        super().__init__(api_key, base_url)

    async def is_available(self) -> bool:
        """Check if OpenAI is available"""
        if not self.api_key or self.api_key == "":
            return False
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{self.base_url}/models",
                    headers={"Authorization": f"Bearer {self.api_key}"},
                    timeout=5.0
                )
                self.available = response.status_code == 200
                return self.available
        except Exception as e:
            logger.error(f"OpenAI availability check failed: {e}")
            return False

    async def chat(
        self,
        messages: List[AIMessage],
        model: str = "gpt-4o-mini",
        temperature: float = 0.7,
        max_tokens: int = 1000,
        stream: bool = False
    ) -> str:
        """Send chat completion to OpenAI"""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.base_url}/chat/completions",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json"
                    },
                    json={
                        "model": model,
                        "messages": [msg.to_dict() for msg in messages],
                        "temperature": temperature,
                        "max_tokens": max_tokens,
                        "stream": stream
                    },
                    timeout=60.0
                )
                response.raise_for_status()
                data = response.json()
                return data["choices"][0]["message"]["content"]
        except Exception as e:
            logger.error(f"OpenAI chat error: {e}")
            raise


class AnthropicProvider(BaseAIProvider):
    """Anthropic API provider (Claude models)"""

    def __init__(self, api_key: str, base_url: str = "https://api.anthropic.com/v1"):
        super().__init__(api_key, base_url)

    async def is_available(self) -> bool:
        """Check if Anthropic is available"""
        if not self.api_key or self.api_key == "":
            return False
        # Anthropic doesn't have a simple health check, assume available if API key exists
        self.available = True
        return True

    async def chat(
        self,
        messages: List[AIMessage],
        model: str = "claude-3-haiku-20240307",
        temperature: float = 0.7,
        max_tokens: int = 1000,
        stream: bool = False
    ) -> str:
        """Send chat completion to Anthropic"""
        try:
            # Convert messages to Anthropic format
            system_msg = None
            converted_messages = []

            for msg in messages:
                if msg.role == "system":
                    system_msg = msg.content
                else:
                    converted_messages.append(msg.to_dict())

            async with httpx.AsyncClient() as client:
                payload = {
                    "model": model,
                    "messages": converted_messages,
                    "max_tokens": max_tokens,
                    "temperature": temperature
                }

                if system_msg:
                    payload["system"] = system_msg

                response = await client.post(
                    f"{self.base_url}/messages",
                    headers={
                        "x-api-key": self.api_key,
                        "anthropic-version": "2023-06-01",
                        "Content-Type": "application/json"
                    },
                    json=payload,
                    timeout=60.0
                )
                response.raise_for_status()
                data = response.json()
                return data["content"][0]["text"]
        except Exception as e:
            logger.error(f"Anthropic chat error: {e}")
            raise


class OllamaProvider(BaseAIProvider):
    """Ollama provider (Local AI, no API key needed!)"""

    def __init__(self, base_url: str = "http://ollama:11434"):
        super().__init__(api_key=None, base_url=base_url)
        self.default_model = "llama3.2:3b"  # Fast and efficient for chatbots

    async def is_available(self) -> bool:
        """Check if Ollama is running"""
        try:
            async with httpx.AsyncClient() as client:
                response = await client.get(
                    f"{self.base_url}/api/tags",
                    timeout=5.0
                )
                self.available = response.status_code == 200
                return self.available
        except Exception as e:
            logger.warning(f"Ollama not available: {e}")
            return False

    async def ensure_model(self, model: str) -> bool:
        """Pull model if not available"""
        try:
            async with httpx.AsyncClient(timeout=300.0) as client:
                # Check if model exists
                response = await client.get(f"{self.base_url}/api/tags")
                models = response.json().get("models", [])
                model_names = [m["name"] for m in models]

                if model not in model_names:
                    logger.info(f"Pulling Ollama model: {model}")
                    # Pull the model
                    async with client.stream(
                        "POST",
                        f"{self.base_url}/api/pull",
                        json={"name": model}
                    ) as resp:
                        async for line in resp.aiter_lines():
                            if line:
                                data = json.loads(line)
                                if data.get("status") == "success":
                                    logger.info(f"Model {model} pulled successfully")
                                    return True
                return True
        except Exception as e:
            logger.error(f"Failed to ensure Ollama model: {e}")
            return False

    async def chat(
        self,
        messages: List[AIMessage],
        model: str = None,
        temperature: float = 0.7,
        max_tokens: int = 1000,
        stream: bool = False
    ) -> str:
        """Send chat completion to Ollama"""
        if model is None:
            model = self.default_model

        # Ensure model is available
        await self.ensure_model(model)

        try:
            # Convert to Ollama format
            prompt = self._convert_to_prompt(messages)

            async with httpx.AsyncClient() as client:
                response = await client.post(
                    f"{self.base_url}/api/generate",
                    json={
                        "model": model,
                        "prompt": prompt,
                        "temperature": temperature,
                        "stream": False
                    },
                    timeout=120.0
                )
                response.raise_for_status()
                data = response.json()
                return data["response"]
        except Exception as e:
            logger.error(f"Ollama chat error: {e}")
            raise

    def _convert_to_prompt(self, messages: List[AIMessage]) -> str:
        """Convert messages to a single prompt"""
        prompt_parts = []
        for msg in messages:
            if msg.role == "system":
                prompt_parts.append(f"System: {msg.content}")
            elif msg.role == "user":
                prompt_parts.append(f"User: {msg.content}")
            elif msg.role == "assistant":
                prompt_parts.append(f"Assistant: {msg.content}")

        prompt_parts.append("Assistant:")
        return "\n\n".join(prompt_parts)


class AIProviderManager:
    """
    Manages multiple AI providers with automatic fallback
    Priority: OpenAI -> Anthropic -> Ollama (Local)
    """

    def __init__(
        self,
        openai_key: Optional[str] = None,
        anthropic_key: Optional[str] = None,
        ollama_url: str = "http://ollama:11434"
    ):
        self.providers: Dict[AIProvider, BaseAIProvider] = {}

        # Initialize providers
        if openai_key:
            self.providers[AIProvider.OPENAI] = OpenAIProvider(openai_key)

        if anthropic_key:
            self.providers[AIProvider.ANTHROPIC] = AnthropicProvider(anthropic_key)

        # Always add Ollama as fallback (no API key needed)
        self.providers[AIProvider.OLLAMA] = OllamaProvider(ollama_url)

        self.priority_order = [AIProvider.OPENAI, AIProvider.ANTHROPIC, AIProvider.OLLAMA]
        self.available_providers: List[AIProvider] = []

    async def initialize(self):
        """Check availability of all providers"""
        self.available_providers = []
        for provider_type in self.priority_order:
            if provider_type in self.providers:
                provider = self.providers[provider_type]
                if await provider.is_available():
                    self.available_providers.append(provider_type)
                    logger.info(f"✓ {provider_type.value} is available")
                else:
                    logger.warning(f"✗ {provider_type.value} is not available")

        if not self.available_providers:
            logger.error("No AI providers available!")
            raise RuntimeError("No AI providers configured or available")

        logger.info(f"Active AI providers: {[p.value for p in self.available_providers]}")

    async def chat(
        self,
        messages: List[AIMessage],
        model: Optional[str] = None,
        temperature: float = 0.7,
        max_tokens: int = 1000,
        preferred_provider: Optional[AIProvider] = None
    ) -> str:
        """
        Send chat request with automatic fallback
        Tries providers in order until one succeeds
        """
        providers_to_try = self.available_providers.copy()

        # If preferred provider specified, try it first
        if preferred_provider and preferred_provider in providers_to_try:
            providers_to_try.remove(preferred_provider)
            providers_to_try.insert(0, preferred_provider)

        last_error = None

        for provider_type in providers_to_try:
            try:
                provider = self.providers[provider_type]
                logger.info(f"Attempting chat with {provider_type.value}")

                # Use provider-specific default model if not specified
                chat_model = model
                if chat_model is None:
                    if provider_type == AIProvider.OPENAI:
                        chat_model = "gpt-4o-mini"
                    elif provider_type == AIProvider.ANTHROPIC:
                        chat_model = "claude-3-haiku-20240307"
                    elif provider_type == AIProvider.OLLAMA:
                        chat_model = "llama3.2:3b"

                response = await provider.chat(
                    messages=messages,
                    model=chat_model,
                    temperature=temperature,
                    max_tokens=max_tokens
                )

                logger.info(f"✓ Successfully got response from {provider_type.value}")
                return response

            except Exception as e:
                last_error = e
                logger.warning(f"✗ {provider_type.value} failed: {e}")
                continue

        # All providers failed
        raise RuntimeError(
            f"All AI providers failed. Last error: {last_error}. "
            f"Available providers: {[p.value for p in self.available_providers]}"
        )

    def get_status(self) -> Dict[str, Any]:
        """Get status of all providers"""
        return {
            "available": [p.value for p in self.available_providers],
            "configured": list(self.providers.keys()),
            "priority_order": [p.value for p in self.priority_order]
        }