"""
Google Gemini LLM provider implementation with integrated tracing.
"""
import os
import logging
import asyncio
from typing import Dict, Any, Optional

from ..core.exceptions import LLMError
from .base import BaseLLMProvider

logger = logging.getLogger(__name__)

class GeminiProvider(BaseLLMProvider):
    """Google Gemini LLM provider implementation with automatic call tracing."""
    
    def __init__(
        self,
        model: str = "gemini-1.5-flash",
        api_key: Optional[str] = None,
        **kwargs
    ):
        """
        Initialize Gemini provider.
        
        Args:
            model: Gemini model name (e.g., "gemini-1.5-flash", "gemini-1.5-pro", "gemini-1.0-pro")
            api_key: Google AI API key
            **kwargs: Additional Gemini-specific parameters
        """
        # Get API key from parameter or environment
        api_key = api_key or os.getenv("GOOGLE_API_KEY") or os.getenv("GEMINI_API_KEY")
        
        super().__init__(model=model, api_key=api_key, **kwargs)
        
        # Gemini-specific default parameters
        self.default_params.update({
            'timeout': kwargs.get('timeout', 60),
            'safety_settings': kwargs.get('safety_settings', None),
            'generation_config': kwargs.get('generation_config', None)
        })
        
        # Lazy-load Gemini client
        self._client = None
    
    @property
    def client(self):
        """Lazy-load Google Generative AI client."""
        if self._client is None:
            try:
                import google.generativeai as genai
                self._validate_api_key()
                
                # Configure the API key
                genai.configure(api_key=self.api_key)
                
                # Create the generative model
                self._client = genai.GenerativeModel(self.model)
                logger.debug("Gemini client initialized")
            except ImportError:
                raise LLMError(
                    "Google Generative AI package not installed. Install with: pip install google-generativeai"
                )
        return self._client
    
    async def _generate_impl(self, prompt: str, **kwargs) -> str:
        """
        Provider-specific implementation of text generation for Gemini.
        
        This method contains the actual Gemini API call logic and is automatically
        wrapped with tracing by the base class generate() method.
        
        Args:
            prompt: Input prompt
            **kwargs: Optional parameters
            
        Returns:
            Generated text response
        """
        try:
            # Merge parameters
            params = self._merge_params(kwargs)
            
            # Prepare generation config
            generation_config = params.get('generation_config', {})
            if not generation_config:
                generation_config = {
                    'max_output_tokens': params.get('max_tokens'),
                    'temperature': params.get('temperature'),
                    'top_p': params.get('top_p')
                }
            
            # Make API call (Gemini's generate_content can be sync or async)
            # For consistency with other providers, we'll run in executor if needed
            if asyncio.iscoroutinefunction(self.client.generate_content):
                response = await self.client.generate_content(
                    prompt,
                    generation_config=generation_config,
                    safety_settings=params.get('safety_settings')
                )
            else:
                # Run synchronous method in executor
                loop = asyncio.get_event_loop()
                response = await loop.run_in_executor(
                    None,
                    lambda: self.client.generate_content(
                        prompt,
                        generation_config=generation_config,
                        safety_settings=params.get('safety_settings')
                    )
                )
            
            # Store usage info if available (Gemini's usage tracking varies)
            if hasattr(response, 'usage_metadata'):
                self._last_usage = response.usage_metadata
            
            return response.text
            
        except Exception as e:
            logger.error(f"Gemini generation failed: {str(e)}")
            raise LLMError(f"Gemini generation failed: {str(e)}")
    
    def _get_last_token_usage(self) -> Dict[str, int]:
        """
        Override base class method to handle Gemini's token format.
        
        Gemini uses different token field names in usage_metadata.
        """
        if self._last_usage:
            # Gemini format varies, try to extract what we can
            prompt_tokens = getattr(self._last_usage, 'prompt_token_count', 0)
            completion_tokens = getattr(self._last_usage, 'candidates_token_count', 0)
            total_tokens = getattr(self._last_usage, 'total_token_count', prompt_tokens + completion_tokens)
            
            return {
                'total_tokens': total_tokens,
                'prompt_tokens': prompt_tokens,
                'completion_tokens': completion_tokens
            }
        
        # Fallback to base class estimation
        return super()._get_last_token_usage()
    
    @property
    def info(self) -> Dict[str, Any]:
        """Get information about the Gemini provider."""
        base_info = super().info
        base_info.update({
            'provider_name': 'Google Gemini',
            'api_compatible': 'Google AI'
        })
        return base_info