"""
Core AI client for Zenoo RPC using LiteLLM.

This module provides the foundation for AI-powered features using LiteLLM
as the unified interface for various LLM providers, with optimized support
for Google Gemini models.
"""

import asyncio
import json
import logging
from typing import Any, Dict, List, Optional, Union, Literal
from dataclasses import dataclass
from enum import Enum

try:
    import litellm
    from litellm import acompletion
    LITELLM_AVAILABLE = True
except ImportError:
    LITELLM_AVAILABLE = False

logger = logging.getLogger(__name__)


class AIProvider(Enum):
    """Supported AI providers."""
    GEMINI = "gemini"
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    AZURE = "azure"


@dataclass
class AIConfig:
    """Configuration for AI client."""
    provider: AIProvider
    model: str
    api_key: str
    base_url: Optional[str] = None
    temperature: float = 0.1
    max_tokens: int = 4096
    timeout: float = 30.0
    max_retries: int = 3
    
    def to_litellm_model(self) -> str:
        """Convert to LiteLLM model format."""
        if self.provider == AIProvider.GEMINI:
            return f"gemini/{self.model}"
        elif self.provider == AIProvider.OPENAI:
            return f"openai/{self.model}"
        elif self.provider == AIProvider.ANTHROPIC:
            return f"anthropic/{self.model}"
        elif self.provider == AIProvider.AZURE:
            return f"azure/{self.model}"
        else:
            return self.model


@dataclass
class AIResponse:
    """Response from AI model."""
    content: str
    model: str
    usage: Dict[str, Any]
    finish_reason: str
    response_time: float


class AIClient:
    """Core AI client using LiteLLM for unified LLM access.
    
    This client provides a consistent interface for interacting with various
    LLM providers through LiteLLM, with optimized configurations for each provider.
    
    Features:
    - Unified interface for multiple LLM providers
    - Async support for high-performance operations
    - Automatic retry logic with exponential backoff
    - Structured output support for Gemini models
    - Comprehensive error handling and logging
    
    Example:
        >>> config = AIConfig(
        ...     provider=AIProvider.GEMINI,
        ...     model="gemini-2.5-flash-lite",
        ...     api_key="your-api-key"
        ... )
        >>> client = AIClient(config)
        >>> await client.initialize()
        >>> 
        >>> response = await client.complete(
        ...     "Explain how to use Odoo RPC",
        ...     system="You are an expert Odoo developer"
        ... )
    """
    
    def __init__(self, config: AIConfig):
        """Initialize AI client with configuration.
        
        Args:
            config: AI configuration including provider, model, and credentials
        """
        if not LITELLM_AVAILABLE:
            raise ImportError(
                "LiteLLM is required for AI features. Install with: pip install litellm"
            )
        
        self.config = config
        self.model = config.to_litellm_model()
        self._initialized = False
        
        # Setup LiteLLM configuration
        self._setup_litellm()
    
    def _setup_litellm(self) -> None:
        """Setup LiteLLM configuration."""
        # Set API key for the provider
        if self.config.provider == AIProvider.GEMINI:
            litellm.api_key = self.config.api_key
        elif self.config.provider == AIProvider.OPENAI:
            litellm.openai_key = self.config.api_key
        elif self.config.provider == AIProvider.ANTHROPIC:
            litellm.anthropic_key = self.config.api_key
        
        # Configure timeouts and retries
        litellm.request_timeout = self.config.timeout
        litellm.num_retries = self.config.max_retries
        
        # Enable logging for debugging
        litellm.set_verbose = logger.isEnabledFor(logging.DEBUG)
    
    async def initialize(self) -> None:
        """Initialize the AI client and validate configuration."""
        try:
            # Set initialized to True temporarily for test
            self._initialized = True

            # Test connection with a simple request
            test_response = await self.complete(
                "Hello",
                max_tokens=10,
                temperature=0
            )

            logger.info(f"AI client initialized successfully with {self.model}")

        except Exception as e:
            # Reset initialized state on failure
            self._initialized = False
            logger.error(f"Failed to initialize AI client: {e}")
            raise
    
    async def complete(
        self,
        prompt: str,
        system: Optional[str] = None,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        response_format: Optional[Dict[str, Any]] = None,
        **kwargs
    ) -> AIResponse:
        """Generate completion using the configured LLM.
        
        Args:
            prompt: User prompt/message
            system: System message for context
            temperature: Sampling temperature (0.0 to 1.0)
            max_tokens: Maximum tokens to generate
            response_format: Response format specification (for structured output)
            **kwargs: Additional parameters for the model
            
        Returns:
            AIResponse with generated content and metadata
            
        Raises:
            RuntimeError: If client is not initialized
            Exception: If completion fails after retries
        """
        if not self._initialized:
            raise RuntimeError("AI client not initialized. Call initialize() first.")
        
        # Prepare messages
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})
        
        # Prepare parameters
        params = {
            "model": self.model,
            "messages": messages,
            "temperature": temperature or self.config.temperature,
            "max_tokens": max_tokens or self.config.max_tokens,
            **kwargs
        }
        
        # Add response format for structured output (Gemini support)
        if response_format and self.config.provider == AIProvider.GEMINI:
            params["response_format"] = response_format
        
        try:
            import time
            start_time = time.time()
            
            # Make async completion request
            response = await acompletion(**params)
            
            response_time = time.time() - start_time
            
            # Extract response data
            choice = response.choices[0]
            content = choice.message.content
            
            return AIResponse(
                content=content,
                model=response.model,
                usage=response.usage.dict() if response.usage else {},
                finish_reason=choice.finish_reason,
                response_time=response_time
            )
            
        except Exception as e:
            logger.error(f"AI completion failed: {e}")
            raise
    
    async def complete_structured(
        self,
        prompt: str,
        schema: Dict[str, Any],
        system: Optional[str] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """Generate structured completion with JSON schema validation.
        
        This method is optimized for Gemini's structured output capabilities
        but provides fallback parsing for other providers.
        
        Args:
            prompt: User prompt/message
            schema: JSON schema for response structure
            system: System message for context
            **kwargs: Additional parameters
            
        Returns:
            Parsed JSON response matching the schema
            
        Raises:
            ValueError: If response doesn't match schema
            json.JSONDecodeError: If response is not valid JSON
        """
        # For Gemini, use native structured output
        if self.config.provider == AIProvider.GEMINI:
            response_format = {"type": "json_object", "schema": schema}
            response = await self.complete(
                prompt=prompt,
                system=system,
                response_format=response_format,
                **kwargs
            )
        else:
            # For other providers, add schema to prompt
            schema_prompt = f"{prompt}\n\nRespond with valid JSON matching this schema:\n{json.dumps(schema, indent=2)}"
            response = await self.complete(
                prompt=schema_prompt,
                system=system,
                **kwargs
            )
        
        try:
            # Parse JSON response
            result = json.loads(response.content)
            return result
            
        except json.JSONDecodeError as e:
            logger.error(f"Failed to parse JSON response: {e}")
            logger.error(f"Response content: {response.content}")
            raise
    
    async def close(self) -> None:
        """Clean up resources."""
        self._initialized = False
        logger.info("AI client closed")
    
    def is_initialized(self) -> bool:
        """Check if client is initialized."""
        return self._initialized
    
    @property
    def provider_name(self) -> str:
        """Get provider name."""
        return self.config.provider.value
    
    @property
    def model_name(self) -> str:
        """Get model name."""
        return self.config.model
