"""
Multi-Provider LLM Client for Toxo integration.

This module provides a unified interface to multiple LLM providers (Gemini, OpenAI, Claude)
with proper error handling, rate limiting, and response processing.
"""

import asyncio
import time
from typing import Optional, Dict, Any, List, Union
from abc import ABC, abstractmethod
import json

from ..core.config import ToxoConfig
from ..utils.logger import get_logger
from ..utils.exceptions import APIError


class BaseLLMClient(ABC):
    """Base class for LLM clients."""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.logger = get_logger(f"toxo.llm.{self.__class__.__name__.lower()}")
    
    @abstractmethod
    async def generate_text(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt."""
        pass
    
    @abstractmethod
    def generate_text_sync(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt (synchronous)."""
        pass
    
    @abstractmethod
    async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for texts."""
        pass


class GeminiClient(BaseLLMClient):
    """Gemini API client."""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        try:
            import google.generativeai as genai
            genai.configure(api_key=config['api_key'])
            self.model = genai.GenerativeModel(config.get('model', 'gemini-2.5-flash-lite'))
        except ImportError:
            raise ImportError("google-generativeai package is required. Install with: pip install google-generativeai")
        except Exception as e:
            raise APIError(f"Failed to initialize Gemini client: {str(e)}")
    
    async def generate_text(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt."""
        # google-generativeai SDK call is synchronous; run in a thread to keep async callers correct.
        return await asyncio.to_thread(self.generate_text_sync, prompt, **kwargs)
    
    def generate_text_sync(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt (synchronous)."""
        try:
            # Use the correct Gemini API format
            response = self.model.generate_content(
                prompt,
                generation_config={
                    "temperature": kwargs.get("temperature", 0.7),
                    "max_output_tokens": kwargs.get("max_tokens", 2048),
                }
            )
            return response.text if response.text else "No response generated"
        except Exception as e:
            raise APIError(f"Gemini API error: {str(e)}")
    
    async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for texts."""
        try:
            import google.generativeai as genai
            result = genai.embed_content(
                model="models/embedding-001",
                content=texts,
                task_type="retrieval_document"
            )
            return result['embedding']
        except Exception as e:
            raise APIError(f"Gemini embeddings error: {str(e)}")


class OpenAIClient(BaseLLMClient):
    """OpenAI API client."""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        try:
            import openai
            self.client = openai.OpenAI(api_key=config['api_key'])
            self.model_name = config.get('model', 'gpt-4')
        except ImportError:
            raise ImportError("openai package is required. Install with: pip install openai")
        except Exception as e:
            raise APIError(f"Failed to initialize OpenAI client: {str(e)}")
    
    async def generate_text(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt."""
        # The OpenAI client used here is synchronous; run in a thread.
        return await asyncio.to_thread(self.generate_text_sync, prompt, **kwargs)
    
    def generate_text_sync(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt (synchronous)."""
        try:
            # Use the correct OpenAI API format with proper message structure
            response = self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": prompt}
                ],
                temperature=kwargs.get("temperature", 0.7),
                max_tokens=kwargs.get("max_tokens", 2048),
                **{k: v for k, v in kwargs.items() if k not in ["temperature", "max_tokens"]}
            )
            return response.choices[0].message.content
        except Exception as e:
            raise APIError(f"OpenAI API error: {str(e)}")
    
    async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for texts."""
        try:
            def _sync():
                response = self.client.embeddings.create(
                    model="text-embedding-3-small",
                    input=texts,
                )
                return [item.embedding for item in response.data]

            return await asyncio.to_thread(_sync)
        except Exception as e:
            raise APIError(f"OpenAI embeddings error: {str(e)}")


class ClaudeClient(BaseLLMClient):
    """Claude API client."""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        try:
            import anthropic
            self.client = anthropic.Anthropic(api_key=config['api_key'])
            self.model_name = config.get('model', 'claude-3.5-sonnet')
        except ImportError:
            raise ImportError("anthropic package is required. Install with: pip install anthropic")
        except Exception as e:
            raise APIError(f"Failed to initialize Claude client: {str(e)}")
    
    async def generate_text(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt."""
        # The anthropic client used here is synchronous; run in a thread.
        return await asyncio.to_thread(self.generate_text_sync, prompt, **kwargs)
    
    def generate_text_sync(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt (synchronous)."""
        try:
            # Use the correct Claude API format
            response = self.client.messages.create(
                model=self.model_name,
                max_tokens=kwargs.get("max_tokens", 300),
                temperature=kwargs.get("temperature", 0.7),
                messages=[{"role": "user", "content": prompt}],
                **{k: v for k, v in kwargs.items() if k not in ["max_tokens", "temperature"]}
            )
            return response.content[0].text
        except Exception as e:
            raise APIError(f"Claude API error: {str(e)}")
    
    async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for texts."""
        try:
            # Claude doesn't have a direct embeddings API, use a workaround
            # For now, return dummy embeddings - in production, you might use a different service
            self.logger.warning("Claude doesn't support embeddings directly. Using dummy embeddings.")
            return [[0.0] * 768 for _ in texts]  # Dummy embeddings
        except Exception as e:
            raise APIError(f"Claude embeddings error: {str(e)}")


class MultiProviderLLMClient:
    """
    Multi-provider LLM client that can switch between different providers.
    
    This class provides a unified interface for multiple LLM providers,
    allowing Toxo layers to be trained and used with any supported provider.
    """
    
    def __init__(self, provider: str = "gemini", config: Optional[Dict[str, Any]] = None):
        """
        Initialize multi-provider LLM client.
        
        Args:
            provider: LLM provider ("gemini", "openai", "claude")
            config: Provider-specific configuration
        """
        self.provider = provider.lower()
        self.config = config or {}
        self.logger = get_logger(f"toxo.llm.{self.provider}")
        
        # Initialize the appropriate client
        if self.provider == "gemini":
            self.client = GeminiClient(self.config)
        elif self.provider == "openai":
            self.client = OpenAIClient(self.config)
        elif self.provider == "claude":
            self.client = ClaudeClient(self.config)
        else:
            raise ValueError(f"Unsupported provider: {provider}. Supported: gemini, openai, claude")
        
        self.logger.info(f"Multi-provider LLM client initialized with {self.provider}")
    
    async def generate_text(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt."""
        return await self.client.generate_text(prompt, **kwargs)
    
    def generate_text_sync(self, prompt: str, **kwargs) -> str:
        """Generate text from prompt (synchronous)."""
        return self.client.generate_text_sync(prompt, **kwargs)

    # Compatibility adapter: Toxo core expects a Gemini-like `.generate(...)`.
    async def generate(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs,
    ) -> str:
        """
        GeminiClient-compatible adapter used by `ToxoLayer`.

        Maps to `generate_text(...)` under the hood.
        """
        if temperature is not None:
            kwargs["temperature"] = temperature
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        return await self.generate_text(prompt, **kwargs)

    def generate_sync(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs,
    ) -> str:
        """Synchronous adapter for `.generate(...)`."""
        if temperature is not None:
            kwargs["temperature"] = temperature
        if max_tokens is not None:
            kwargs["max_tokens"] = max_tokens
        return self.generate_text_sync(prompt, **kwargs)
    
    async def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
        """Generate embeddings for texts."""
        return await self.client.generate_embeddings(texts)
    
    def switch_provider(self, provider: str, config: Dict[str, Any]):
        """Switch to a different provider."""
        self.provider = provider.lower()
        self.config = config
        
        if self.provider == "gemini":
            self.client = GeminiClient(self.config)
        elif self.provider == "openai":
            self.client = OpenAIClient(self.config)
        elif self.provider == "claude":
            self.client = ClaudeClient(self.config)
        else:
            raise ValueError(f"Unsupported provider: {provider}")
        
        self.logger.info(f"Switched to {self.provider} provider")
    
    def get_provider_info(self) -> Dict[str, Any]:
        """Get information about the current provider."""
        return {
            "provider": self.provider,
            "model": self.config.get('model', 'unknown'),
            "config": self.config
        }


def create_llm_client(provider: str, api_key: str, model: str) -> MultiProviderLLMClient:
    """
    Factory function to create an LLM client.
    
    Args:
        provider: LLM provider ("gemini", "openai", "claude")
        api_key: API key for the provider
        model: Model name (user must specify)
    
    Returns:
        MultiProviderLLMClient instance
    """
    if not model:
        raise ValueError("Model name is required. User must specify the model.")
    
    config = {
        "api_key": api_key,
        "model": model
    }
    
    return MultiProviderLLMClient(provider, config)
