"""
Gemini API client for Toxo integration.

This module provides a clean interface to Google's Gemini API with
proper error handling, rate limiting, and response processing.
"""

import asyncio
import time
import base64
from typing import Optional, Dict, Any, List, Union
from pathlib import Path
from io import BytesIO
try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    Image = None
    PIL_AVAILABLE = False
import json
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold

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


class GeminiClient:
    """
    Client for Google Gemini API integration.
    
    Provides async/sync interfaces for text generation, embeddings,
    vision processing, and other Gemini capabilities with proper error handling.
    """
    
    def __init__(self, config: GeminiConfig):
        """
        Initialize Gemini client.
        
        Args:
            config: Gemini configuration
        """
        self.config = config
        self.logger = get_logger("toxo.gemini")
        
        # Configure the API
        genai.configure(api_key=config.api_key)
        # Never log API keys (even partially).
        self.logger.info(f"GeminiClient initialized (model={config.model})")
        
        # Initialize models
        self.text_model = genai.GenerativeModel(
            model_name=config.model,
            safety_settings={
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            }
        )
        
        # Initialize vision model for multimodal processing
        vision_model_name = getattr(config, "vision_model", None) or config.model
        self.vision_model = genai.GenerativeModel(
            model_name=vision_model_name,
            safety_settings={
                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
            }
        )
        
        # Rate limiting and timeout configuration
        self.last_request_time = 0
        self.min_request_interval = 0.1  # Minimum time between requests
        self.default_timeout = 300  # 5 minutes default timeout
        self.max_retries = 3
        self.retry_delays = [1, 5, 15]  # Exponential backoff delays
        
        # Vision processing capabilities
        self.vision_capabilities = {
            'document_analysis': True,
            'table_extraction': True,
            'chart_interpretation': True,
            'handwriting_recognition': True,
            'diagram_understanding': True,
            'layout_analysis': True,
            'formula_recognition': True
        }
        
        self.logger.info(f"Gemini client initialized with models: {config.model} (text), {vision_model_name} (vision)")

    async def _execute_with_timeout_and_retry(
        self,
        operation,
        operation_name: str,
        timeout: int = None,
        max_retries: int = None,
        *args,
        **kwargs
    ) -> Any:
        """
        Execute an operation with timeout and retry logic.
        
        Args:
            operation: The operation to execute
            operation_name: Name for logging
            timeout: Timeout in seconds
            max_retries: Maximum retry attempts
            *args: Positional arguments for the operation
            **kwargs: Keyword arguments for the operation
            
        Returns:
            Operation result
        """
        timeout = timeout or self.default_timeout
        max_retries = max_retries or self.max_retries
        
        for attempt in range(max_retries + 1):
            try:
                # Execute with timeout
                result = await asyncio.wait_for(
                    asyncio.to_thread(operation, *args, **kwargs),
                    timeout=timeout
                )
                return result
                
            except asyncio.TimeoutError:
                self.logger.warning(f"{operation_name} timed out after {timeout}s (attempt {attempt + 1}/{max_retries + 1})")
                if attempt < max_retries:
                    await asyncio.sleep(self.retry_delays[min(attempt, len(self.retry_delays) - 1)])
                    continue
                raise APIError(f"{operation_name} failed: Timeout after {timeout}s")
                
            except Exception as e:
                error_msg = str(e)
                if "504" in error_msg or "deadline" in error_msg.lower():
                    self.logger.warning(f"{operation_name} got 504/deadline error (attempt {attempt + 1}/{max_retries + 1}): {error_msg}")
                    if attempt < max_retries:
                        await asyncio.sleep(self.retry_delays[min(attempt, len(self.retry_delays) - 1)])
                        continue
                    raise APIError(f"{operation_name} failed: {error_msg}")
                else:
                    # Non-retryable error
                    self.logger.error(f"{operation_name} failed with non-retryable error: {error_msg}")
                    raise APIError(f"{operation_name} failed: {error_msg}")
        
        raise APIError(f"{operation_name} failed after {max_retries + 1} attempts")

    async def generate(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        timeout: Optional[int] = None,
        **kwargs
    ) -> str:
        """
        Generate text using Gemini.
        
        Args:
            prompt: Input prompt
            temperature: Generation temperature (0-1)
            max_tokens: Maximum tokens to generate
            timeout: Request timeout in seconds
            **kwargs: Additional generation parameters
            
        Returns:
            Generated text
        """
        await self._rate_limit()
        
        # Use config defaults if not specified
        temperature = temperature or self.config.temperature
        max_tokens = max_tokens or self.config.max_tokens
        timeout = timeout or (240 if len(prompt) > 5000 else 120)  # Adaptive timeout
        
        try:
            # Configure generation parameters
            generation_config = genai.types.GenerationConfig(
                temperature=temperature,
                max_output_tokens=max_tokens,
                **kwargs
            )
            
            # Execute with timeout and retry - use correct parameter name
            response = await self._execute_with_timeout_and_retry(
                self.text_model.generate_content,
                "Text generation",
                timeout,
                None,  # use default max_retries
                prompt,  # positional argument
                generation_config=generation_config
            )
            
            # Extract text from response
            if response.candidates and response.candidates[0].content:
                generated_text = response.candidates[0].content.parts[0].text
                self.logger.debug(f"Generated {len(generated_text)} characters")
                return generated_text
            else:
                self.logger.warning("Empty response from Gemini")
                return ""
                
        except Exception as e:
            self.logger.error(f"Gemini generation failed: {str(e)}")
            raise APIError(f"Generation failed: {str(e)}")

    def generate_sync(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> str:
        """
        Synchronous version of generate.
        
        Args:
            prompt: Input prompt
            temperature: Generation temperature
            max_tokens: Maximum tokens
            **kwargs: Additional parameters
            
        Returns:
            Generated text
        """
        return asyncio.run(self.generate(prompt, temperature, max_tokens, **kwargs))

    async def analyze_image(
        self,
        image_data: Union[bytes, str, Path, Any],
        prompt: str = "Analyze this image in detail.",
        analysis_type: str = "comprehensive",
        timeout: Optional[int] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Analyze image using Gemini Vision.
        
        Args:
            image_data: Image data (bytes, base64, PIL Image, or file path)
            prompt: Analysis prompt
            analysis_type: Type of analysis (comprehensive, document, chart, etc.)
            timeout: Request timeout in seconds
            **kwargs: Additional parameters
            
        Returns:
            Analysis results
        """
        await self._rate_limit()
        
        try:
            # Process image data
            image_part = await self._prepare_image_part(image_data)
            
            # Create specialized prompt based on analysis type
            specialized_prompt = self._create_vision_prompt(prompt, analysis_type)
            
            # Adaptive timeout based on analysis type
            timeout = timeout or (180 if analysis_type in ["comprehensive", "document"] else 90)
            
            # Generate response with image - use correct parameter format
            response = await self._execute_with_timeout_and_retry(
                self.vision_model.generate_content,
                "Image analysis",
                timeout,
                None,  # use default max_retries
                [specialized_prompt, image_part],  # positional argument
                generation_config=genai.types.GenerationConfig(
                    temperature=kwargs.get('temperature', 0.2),  # Lower temp for factual analysis
                    max_output_tokens=kwargs.get('max_tokens', 2048)
                )
            )
            
            if response.candidates and response.candidates[0].content:
                analysis_text = response.candidates[0].content.parts[0].text
                
                # Parse structured response if possible
                try:
                    structured_analysis = self._parse_vision_response(analysis_text, analysis_type)
                    return {
                        'analysis_type': analysis_type,
                        'raw_response': analysis_text,
                        'structured_analysis': structured_analysis,
                        'metadata': {
                            'model': 'gemini-1.5-flash',
                            'timestamp': time.time(),
                            'prompt_type': analysis_type
                        }
                    }
                except Exception as parse_error:
                    self.logger.debug(f"Could not parse structured response: {parse_error}")
                    return {
                        'analysis_type': analysis_type,
                        'raw_response': analysis_text,
                        'structured_analysis': None,
                        'metadata': {
                            'model': 'gemini-1.5-flash',
                            'timestamp': time.time(),
                            'prompt_type': analysis_type
                        }
                    }
            else:
                self.logger.warning("Empty response from Gemini Vision")
                return {
                    'analysis_type': analysis_type,
                    'raw_response': "",
                    'structured_analysis': None,
                    'error': 'Empty response'
                }
                
        except Exception as e:
            self.logger.error(f"Image analysis failed: {str(e)}")
            raise APIError(f"Image analysis failed: {str(e)}")

    async def analyze_document(
        self,
        document_data: Union[bytes, str, Path, Any, List[Union[bytes, str, Path, Any]]],
        document_type: str = "auto",
        extract_tables: bool = True,
        extract_charts: bool = True,
        extract_text: bool = True,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Comprehensive document analysis using Gemini Vision.
        
        Args:
            document_data: Document data (can be multiple pages/images)
            document_type: Type of document (pdf, report, form, invoice, etc.)
            extract_tables: Whether to extract table data
            extract_charts: Whether to interpret charts/graphs
            extract_text: Whether to extract all text content
            **kwargs: Additional parameters
            
        Returns:
            Comprehensive document analysis
        """
        await self._rate_limit()
        
        try:
            # Handle multiple pages/images
            if isinstance(document_data, list):
                pages_analysis = []
                for i, page_data in enumerate(document_data):
                    page_analysis = await self._analyze_single_page(
                        page_data, document_type, extract_tables, extract_charts, extract_text, i + 1, **kwargs
                    )
                    pages_analysis.append(page_analysis)
                
                # Combine analysis from all pages
                combined_analysis = self._combine_document_analysis(pages_analysis, document_type)
                return combined_analysis
            else:
                # Single page/image analysis
                return await self._analyze_single_page(
                    document_data, document_type, extract_tables, extract_charts, extract_text, 1, **kwargs
                )
                
        except Exception as e:
            self.logger.error(f"Document analysis failed: {str(e)}")
            raise APIError(f"Document analysis failed: {str(e)}")

    async def generate_with_file(
        self,
        *,
        prompt: str,
        file_bytes: bytes,
        mime_type: str,
        analysis_type: str = "document",
        timeout: Optional[int] = None,
        **kwargs,
    ) -> str:
        """
        Generate a response from Gemini given a binary attachment (PDF/image/etc).

        This is the lowest-dependency path for multimodal: we send bytes directly to Gemini
        and do not require local PDF rendering / OCR.
        """
        await self._rate_limit()

        try:
            specialized_prompt = self._create_vision_prompt(prompt, analysis_type)
            timeout = timeout or (180 if analysis_type in ["comprehensive", "document"] else 90)
            part = {"mime_type": mime_type, "data": file_bytes}

            gen_cfg = genai.types.GenerationConfig(
                temperature=kwargs.get("temperature", 0.2),
                max_output_tokens=kwargs.get("max_tokens", kwargs.get("max_output_tokens", 2048)),
            )

            # Prefer the vision model, but fall back to the main text model if the
            # vision model is unavailable for this API version/account.
            try:
                response = await self._execute_with_timeout_and_retry(
                    self.vision_model.generate_content,
                    "File-based generation",
                    timeout,
                    None,
                    [specialized_prompt, part],
                    generation_config=gen_cfg,
                )
            except Exception as e:
                msg = str(e).lower()
                if "not found" in msg or "is not supported" in msg:
                    response = await self._execute_with_timeout_and_retry(
                        self.text_model.generate_content,
                        "File-based generation (fallback)",
                        timeout,
                        None,
                        [specialized_prompt, part],
                        generation_config=gen_cfg,
                    )
                else:
                    raise

            if response.candidates and response.candidates[0].content:
                return response.candidates[0].content.parts[0].text
            return ""
        except Exception as e:
            self.logger.error(f"File-based generation failed: {str(e)}")
            raise APIError(f"File-based generation failed: {str(e)}")

    async def analyze_dataset(
        self,
        dataset_info: Dict[str, Any],
        sample_data: Optional[str] = None,
        analysis_goals: List[str] = None,
        timeout: Optional[int] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Analyze dataset using Gemini with enhanced error handling.
        
        Args:
            dataset_info: Information about the dataset
            sample_data: Sample data for analysis
            analysis_goals: Specific analysis goals
            timeout: Request timeout in seconds
            **kwargs: Additional parameters
            
        Returns:
            Analysis results
        """
        await self._rate_limit()
        
        try:
            # Set adaptive timeout based on dataset complexity
            timeout = timeout or (300 if len(str(dataset_info)) > 2000 else 180)
            
            # Create analysis prompt
            analysis_prompt = self._create_dataset_analysis_prompt(
                dataset_info, sample_data, analysis_goals or ['insights', 'quality', 'recommendations']
            )
            
            # Limit prompt size to prevent timeout
            if len(analysis_prompt) > 8000:
                analysis_prompt = analysis_prompt[:8000] + "\n\n... (truncated for analysis)"
            
            # Generate analysis with timeout and retry - use correct parameter format
            response = await self._execute_with_timeout_and_retry(
                self.text_model.generate_content,
                "Dataset analysis",
                timeout,
                None,  # use default max_retries
                analysis_prompt,  # positional argument
                generation_config=genai.types.GenerationConfig(
                temperature=kwargs.get('temperature', 0.3),
                    max_output_tokens=kwargs.get('max_tokens', 2048)
                )
            )
            
            if response.candidates and response.candidates[0].content:
                analysis_text = response.candidates[0].content.parts[0].text
                
                # Parse structured response
                try:
                    parsed_analysis = self._parse_dataset_response(analysis_text)
                    if parsed_analysis:
                        return {
                            'analysis_type': 'dataset',
                            'raw_response': analysis_text,
                            'structured_analysis': parsed_analysis,
                            'dataset_info': dataset_info,
                            'analysis_timestamp': time.time()
                        }
                except Exception as parse_error:
                    self.logger.warning(f"Failed to parse dataset analysis: {parse_error}")
                
                # Return raw response if parsing fails
                return {
                    'analysis_type': 'dataset',
                    'raw_response': analysis_text,
                    'dataset_info': dataset_info,
                    'analysis_timestamp': time.time()
                }
            else:
                self.logger.warning("Empty response from dataset analysis")
                return {
                    'analysis_type': 'dataset',
                    'error': 'Empty response from Gemini',
                    'dataset_info': dataset_info
                }
                
        except Exception as e:
            self.logger.error(f"Dataset analysis failed: {str(e)}")
            raise APIError(f"Dataset analysis failed: {str(e)}")

    async def analyze_content(
        self,
        content: str,
        analysis_type: str = "sentiment",
        timeout: Optional[int] = None,
        **kwargs
    ) -> Dict[str, Any]:
        """
        Analyze content with Gemini using robust error handling.
        
        Args:
            content: Content to analyze
            analysis_type: Type of analysis
            timeout: Request timeout in seconds
            **kwargs: Additional parameters
            
        Returns:
            Analysis results
        """
        await self._rate_limit()
        
        try:
            # Set adaptive timeout
            timeout = timeout or (180 if len(content) > 3000 else 90)
            
            # Create analysis prompts based on type
            if analysis_type == "comprehensive":
                prompt = f"""Analyze this content comprehensively and provide insights:

{content[:6000]}

Provide a detailed analysis including:
1. Document type and purpose
2. Key information and entities
3. Structure and organization
4. Quality assessment
5. Business context and implications
6. Actionable insights and recommendations

Format your response as structured text with clear sections."""

            elif analysis_type == "sentiment":
                prompt = f"""Analyze the sentiment of this content:

{content[:3000]}

Provide:
1. Overall sentiment (positive/negative/neutral)
2. Confidence score (0-1)
3. Key emotional indicators
4. Sentiment breakdown by section if applicable"""

            elif analysis_type == "keywords":
                prompt = f"""Extract key terms and topics from this content:

{content[:4000]}

Provide:
1. Main keywords and phrases
2. Topics and themes
3. Important entities (people, places, organizations)
4. Technical terms or jargon"""

            elif analysis_type == "summary":
                prompt = f"""Summarize this content concisely:

{content[:5000]}

Provide:
1. Brief summary (2-3 sentences)
2. Key points (bullet list)
3. Main conclusions
4. Important details"""

            elif analysis_type == "classification":
                prompt = f"""Classify and categorize this content:

{content[:4000]}

Determine:
1. Document type/category
2. Industry or domain
3. Purpose and audience
4. Content structure
5. Classification confidence (0-1)"""

            else:
                prompt = f"""Analyze this content for {analysis_type}:

{content[:4000]}

Provide relevant insights and analysis."""

            # Generate analysis with timeout and retry - use correct parameter format
            response = await self._execute_with_timeout_and_retry(
                self.text_model.generate_content,
                f"Content analysis ({analysis_type})",
                timeout,
                None,  # use default max_retries
                prompt,  # positional argument
                generation_config=genai.types.GenerationConfig(
                    temperature=kwargs.get('temperature', 0.3),
                    max_output_tokens=kwargs.get('max_tokens', 1500)
                )
            )
            
            if response.candidates and response.candidates[0].content:
                analysis_text = response.candidates[0].content.parts[0].text
                
                return {
                    'analysis_type': analysis_type,
                    'content_length': len(content),
                    'analysis_result': analysis_text,
                    'timestamp': time.time()
                }
            else:
                self.logger.warning(f"Empty response from content analysis ({analysis_type})")
                return {
                    'analysis_type': analysis_type,
                    'error': 'Empty response from Gemini',
                    'content_length': len(content)
                }
                
        except Exception as e:
            self.logger.error(f"Content analysis failed: {str(e)}")
            raise APIError(f"Analysis failed: {str(e)}")

    async def extract_structured_data(
        self,
        image_data: Union[bytes, str, Path, Any],
        data_type: str = "table",
        output_format: str = "json",
        **kwargs
    ) -> Dict[str, Any]:
        """
        Extract structured data from images (tables, forms, etc.).
        
        Args:
            image_data: Image containing structured data
            data_type: Type of data to extract (table, form, invoice, etc.)
            output_format: Output format (json, csv, etc.)
            **kwargs: Additional parameters
            
        Returns:
            Extracted structured data
        """
        await self._rate_limit()
        
        try:
            image_part = await self._prepare_image_part(image_data)
            
            # Create extraction prompt
            extraction_prompt = self._create_extraction_prompt(data_type, output_format)
            
            response = await asyncio.to_thread(
                self.vision_model.generate_content,
                [extraction_prompt, image_part],
                generation_config=genai.types.GenerationConfig(
                    temperature=0.1,  # Very low temperature for accuracy
                    max_output_tokens=kwargs.get('max_tokens', 2048)
                )
            )
            
            if response.candidates and response.candidates[0].content:
                extraction_text = response.candidates[0].content.parts[0].text
                
                # Parse extracted data
                structured_data = self._parse_extracted_data(extraction_text, data_type, output_format)
                
                return {
                    'data_type': data_type,
                    'output_format': output_format,
                    'raw_extraction': extraction_text,
                    'structured_data': structured_data,
                    'metadata': {
                        'model': 'gemini-1.5-flash',
                        'timestamp': time.time(),
                        'extraction_type': data_type
                    }
                }
            else:
                return {
                    'data_type': data_type,
                    'output_format': output_format,
                    'raw_extraction': "",
                    'structured_data': None,
                    'error': 'Empty response'
                }
                
        except Exception as e:
            self.logger.error(f"Structured data extraction failed: {str(e)}")
            raise APIError(f"Structured data extraction failed: {str(e)}")

    async def embed_text(self, text: str) -> List[float]:
        """
        Get embeddings for text using Gemini.
        
        Args:
            text: Text to embed
            
        Returns:
            Embedding vector
        """
        await self._rate_limit()
        
        try:
            # Embedding models vary by API version/account. We:
            # - try a short fallback chain
            # - then dynamically look up a supported embed model via list_models()
            candidate_models = [
                "models/text-embedding-004",
                "models/embedding-001",
            ]

            def _try_embed(model_name: str):
                return genai.embed_content(model=model_name, content=text)

            last_error: Optional[Exception] = None
            response = None
            for embedding_model in candidate_models:
                try:
                    response = await asyncio.to_thread(_try_embed, embedding_model)
                    last_error = None
                    break
                except Exception as e:
                    last_error = e

            if response is None:
                # Dynamic discovery fallback (one-time per process).
                discovered = getattr(self, "_embedding_model_name", None)
                if not discovered:
                    try:
                        models = await asyncio.to_thread(genai.list_models)
                        for m in models:
                            name = getattr(m, "name", None)
                            methods = getattr(m, "supported_generation_methods", None) or []
                            if name and ("embedContent" in methods or "embedText" in methods):
                                discovered = name if name.startswith("models/") else f"models/{name}"
                                break
                    except Exception as e:
                        last_error = e

                    if discovered:
                        setattr(self, "_embedding_model_name", discovered)

                if discovered:
                    response = await asyncio.to_thread(_try_embed, discovered)
                else:
                    raise last_error or RuntimeError("No embedding model succeeded")

            embedding = response["embedding"]
            self.logger.debug(f"Generated embedding of dimension {len(embedding)}")
            return embedding
            
        except Exception as e:
            self.logger.error(f"Embedding generation failed: {str(e)}")
            raise APIError(f"Embedding failed: {str(e)}")

    async def chat(
        self,
        messages: List[Dict[str, str]],
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> str:
        """
        Chat with Gemini using conversation history.
        
        Args:
            messages: List of message dictionaries with 'role' and 'content'
            temperature: Generation temperature
            max_tokens: Maximum tokens
            **kwargs: Additional parameters
            
        Returns:
            Assistant's response
        """
        await self._rate_limit()
        
        # Convert messages to Gemini format
        chat_history = []
        for msg in messages[:-1]:  # All but the last message
            if msg['role'] == 'user':
                chat_history.append({'role': 'user', 'parts': [msg['content']]})
            elif msg['role'] == 'assistant':
                chat_history.append({'role': 'model', 'parts': [msg['content']]})
        
        # Get the current user message
        current_message = messages[-1]['content'] if messages else ""
        
        try:
            # Start chat with history
            chat = self.text_model.start_chat(history=chat_history)
            
            # Configure generation
            generation_config = genai.types.GenerationConfig(
                temperature=temperature or self.config.temperature,
                max_output_tokens=max_tokens or self.config.max_tokens,
                **kwargs
            )
            
            # Send message
            response = await asyncio.to_thread(
                chat.send_message,
                current_message,
                generation_config=generation_config
            )
            
            if response.candidates and response.candidates[0].content:
                response_text = response.candidates[0].content.parts[0].text
                self.logger.debug(f"Chat response: {len(response_text)} characters")
                return response_text
            else:
                self.logger.warning("Empty chat response from Gemini")
                return ""
                
        except Exception as e:
            self.logger.error(f"Chat failed: {str(e)}")
            raise APIError(f"Chat failed: {str(e)}")

    # Helper methods
    async def _prepare_image_part(self, image_data: Union[bytes, str, Path, Any]) -> Any:
        """Prepare image data for Gemini Vision API."""
        if isinstance(image_data, Path):
            # Read file
            with open(image_data, 'rb') as f:
                image_bytes = f.read()
            return {"mime_type": "image/png", "data": image_bytes}
        elif PIL_AVAILABLE and isinstance(image_data, Image.Image):
            # Convert PIL Image to bytes
            byte_io = BytesIO()
            image_data.save(byte_io, format='PNG')
            image_bytes = byte_io.getvalue()
            return {"mime_type": "image/png", "data": image_bytes}
        elif isinstance(image_data, str):
            # Assume base64 encoded
            image_bytes = base64.b64decode(image_data)
            return {"mime_type": "image/png", "data": image_bytes}
        elif isinstance(image_data, bytes):
            # Raw bytes
            return {"mime_type": "image/png", "data": image_data}
        else:
            raise ValueError(f"Unsupported image data type: {type(image_data)}")

    def _create_vision_prompt(self, base_prompt: str, analysis_type: str) -> str:
        """Create specialized vision analysis prompt."""
        prompts = {
            "comprehensive": f"""
            {base_prompt}
            
            Please provide a comprehensive analysis including:
            1. Overall description and content type
            2. Key visual elements (text, images, charts, tables)
            3. Document structure and layout
            4. Data extraction opportunities
            5. Quality assessment
            6. Actionable insights
            
            Format your response in clear sections.
            """,
            "document": f"""
            {base_prompt}
            
            Analyze this document focusing on:
            1. Document type and purpose
            2. Text content and structure
            3. Tables and their data
            4. Charts/graphs and their insights
            5. Key information extraction
            6. Document quality and readability
            
            Provide structured output where possible.
            """,
            "table": f"""
            {base_prompt}
            
            Focus on extracting and analyzing tabular data:
            1. Identify all tables in the image
            2. Extract table headers and data
            3. Analyze data patterns and insights
            4. Provide structured JSON format if possible
            5. Note any data quality issues
            """,
            "chart": f"""
            {base_prompt}
            
            Analyze charts and graphs in this image:
            1. Chart type (bar, line, pie, etc.)
            2. Data being presented
            3. Key insights and trends
            4. Data values and patterns
            5. Recommendations based on the data
            """
        }
        
        return prompts.get(analysis_type, base_prompt)

    def _create_extraction_prompt(self, data_type: str, output_format: str) -> str:
        """Create data extraction prompt."""
        prompts = {
            "table": f"""
            Extract all tabular data from this image. 
            
            Instructions:
            1. Identify all tables
            2. Extract headers and all data rows
            3. Maintain data relationships and structure
            4. Output in {output_format} format
            5. If multiple tables, separate them clearly
            
            Be precise with data extraction - accuracy is critical.
            """,
            "form": f"""
            Extract all form fields and their values from this image.
            
            Instructions:
            1. Identify all form fields (labels and values)
            2. Extract field names and corresponding values
            3. Note any checkboxes or selections
            4. Output in {output_format} format
            5. Preserve field relationships
            
            Be precise with data extraction.
            """,
            "invoice": f"""
            Extract all invoice information from this image.
            
            Instructions:
            1. Extract invoice number, date, amounts
            2. Extract vendor and customer information
            3. Extract line items with quantities and prices
            4. Calculate totals and verify accuracy
            5. Output in {output_format} format
            
            Focus on financial accuracy.
            """
        }
        
        return prompts.get(data_type, f"Extract structured data from this image in {output_format} format.")

    def _create_dataset_analysis_prompt(
        self,
        dataset_info: Dict[str, Any],
        sample_data: Optional[str],
        analysis_goals: List[str]
    ) -> str:
        """Create dataset analysis prompt."""
        prompt = f"""
        Analyze this dataset and provide comprehensive insights:
        
        Dataset Information:
        {json.dumps(dataset_info, indent=2)}
        
        Sample Data:
        {sample_data or 'No sample data provided'}
        
        Analysis Goals:
        {', '.join(analysis_goals) if analysis_goals else 'General analysis'}
        
        Please provide:
        1. Dataset overview and structure assessment
        2. Data quality evaluation
        3. Key patterns and insights
        4. Potential analysis opportunities
        5. Data preprocessing recommendations
        6. Modeling suggestions
        7. Business value assessment
        8. Risk assessment and limitations
        
        Format your response as structured analysis with clear sections.
        """
        
        return prompt

    async def _analyze_single_page(
        self,
        page_data: Union[bytes, str, Path, Any],
        document_type: str,
        extract_tables: bool,
        extract_charts: bool,
        extract_text: bool,
        page_number: int,
        **kwargs
    ) -> Dict[str, Any]:
        """Analyze a single page/image of a document."""
        # Determine analysis focus
        analysis_components = []
        if extract_text:
            analysis_components.append("text content")
        if extract_tables:
            analysis_components.append("tables and structured data")
        if extract_charts:
            analysis_components.append("charts and visualizations")
        
        analysis_prompt = f"""
        Analyze page {page_number} of this {document_type} document.
        
        Focus on extracting: {', '.join(analysis_components)}
        
        Provide structured analysis including:
        1. Page content overview
        2. Text content (if requested)
        3. Tables with data extraction (if requested)
        4. Charts/graphs with insights (if requested)
        5. Key information and insights
        6. Data quality assessment
        
        Format output for easy parsing and integration.
        """
        
        analysis_result = await self.analyze_image(
            page_data,
            analysis_prompt,
            analysis_type="document",
            **kwargs
        )
        
        analysis_result['page_number'] = page_number
        return analysis_result

    def _combine_document_analysis(self, pages_analysis: List[Dict[str, Any]], document_type: str) -> Dict[str, Any]:
        """Combine analysis from multiple pages into comprehensive document analysis."""
        combined = {
            'document_type': document_type,
            'total_pages': len(pages_analysis),
            'pages': pages_analysis,
            'combined_insights': {
                'all_text': "",
                'all_tables': [],
                'all_charts': [],
                'key_insights': []
            },
            'metadata': {
                'analysis_timestamp': time.time(),
                'model': 'gemini-1.5-flash',
                'processing_type': 'multi_page_document'
            }
        }
        
        # Aggregate content from all pages
        for page in pages_analysis:
            if page.get('raw_response'):
                combined['combined_insights']['all_text'] += f"\n--- Page {page.get('page_number', '?')} ---\n"
                combined['combined_insights']['all_text'] += page['raw_response']
        
        return combined

    def _parse_vision_response(self, response_text: str, analysis_type: str) -> Optional[Dict[str, Any]]:
        """Parse structured response from vision analysis."""
        # This is a simplified parser - in production, you'd want more sophisticated parsing
        try:
            # Look for JSON-like structures
            if '{' in response_text and '}' in response_text:
                # Try to extract JSON
                start = response_text.find('{')
                end = response_text.rfind('}') + 1
                json_str = response_text[start:end]
                return json.loads(json_str)
            
            # For now, return basic structure
            return {
                'analysis_type': analysis_type,
                'parsed_content': response_text,
                'structure_detected': 'text_only'
            }
        except Exception:
            return None

    def _parse_dataset_response(self, response_text: str) -> Optional[Dict[str, Any]]:
        """Parse dataset analysis response."""
        # Simplified parsing - in production, implement more sophisticated parsing
        return {
            'raw_analysis': response_text,
            'parsed_sections': response_text.split('\n\n'),
            'analysis_timestamp': time.time()
        }

    def _parse_extracted_data(self, extraction_text: str, data_type: str, output_format: str) -> Optional[Any]:
        """Parse extracted structured data."""
        try:
            if output_format.lower() == 'json':
                # Try to find and parse JSON
                if '{' in extraction_text:
                    start = extraction_text.find('{')
                    end = extraction_text.rfind('}') + 1
                    json_str = extraction_text[start:end]
                    return json.loads(json_str)
            
            # For other formats, return as-is for now
            return {
                'raw_data': extraction_text,
                'data_type': data_type,
                'format': output_format
            }
        except Exception as e:
            self.logger.debug(f"Could not parse extracted data: {e}")
            return None

    async def _rate_limit(self):
        """Implement basic rate limiting."""
        current_time = time.time()
        time_since_last = current_time - self.last_request_time
        
        if time_since_last < self.min_request_interval:
            wait_time = self.min_request_interval - time_since_last
            await asyncio.sleep(wait_time)
        
        self.last_request_time = time.time()

    def get_model_info(self) -> Dict[str, Any]:
        """Get information about the current model."""
        try:
            model_info = genai.get_model(f"models/{self.config.model}")
            return {
                "name": model_info.name,
                "display_name": model_info.display_name,
                "description": model_info.description,
                "input_token_limit": model_info.input_token_limit,
                "output_token_limit": model_info.output_token_limit,
                "supported_generation_methods": model_info.supported_generation_methods,
                "vision_capabilities": self.vision_capabilities
            }
        except Exception as e:
            self.logger.warning(f"Could not get model info: {str(e)}")
            return {"name": self.config.model, "error": str(e)}

    def test_connection(self) -> bool:
        """Test connection to Gemini API."""
        try:
            # Simple test generation
            response = self.text_model.generate_content("Hello")
            return bool(response.candidates and response.candidates[0].content)
        except Exception as e:
            self.logger.error(f"Connection test failed: {str(e)}")
            return False

    async def generate_content(
        self,
        prompt: str,
        temperature: Optional[float] = None,
        max_tokens: Optional[int] = None,
        timeout: Optional[int] = None,
        **kwargs
    ) -> str:
        """
        Generate content using Gemini (alias for generate method).
        
        Args:
            prompt: Input prompt
            temperature: Generation temperature (0-1)
            max_tokens: Maximum tokens to generate
            timeout: Request timeout in seconds
            **kwargs: Additional generation parameters
            
        Returns:
            Generated text
        """
        return await self.generate(
            prompt=prompt,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout=timeout,
            **kwargs
        ) 