# llm_clients/gemini.py
from typing import Generator, Any, Dict, List, Optional
import google.generativeai as genai
from .base import LlmClient
import config
import tempfile
import os

class GeminiClient(LlmClient):
    """LLM client for Google's Gemini models."""

    def __init__(self, config_dict: Dict[str, Any], system_prompt: str):
        super().__init__(config_dict, system_prompt)
        if not config.GEMINI_API_KEY or "YOUR_GOOGLE_API_KEY" in config.GEMINI_API_KEY:
            raise ValueError("Please set your GEMINI_API_KEY in the config.py file or as an environment variable.")
        
        genai.configure(api_key=config.GEMINI_API_KEY)
        self.model = genai.GenerativeModel(
            self.config['model'],
            system_instruction=self.system_prompt
        )
        print(f"✅ Gemini Client initialized with model '{self.config['model']}'.")

    def generate_content(self, prompt: str, files: Optional[List[Dict[str, Any]]] = None) -> str:
        """Generates a non-streaming response from Gemini with optional file attachments."""
        content_parts = [prompt]
        
        # Upload files to Gemini if provided
        if files:
            for file_info in files:
                try:
                    # Write file content to temporary file
                    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_info['extension'].lstrip('.')}") as tmp_file:
                        tmp_file.write(file_info['content'])
                        tmp_file_path = tmp_file.name
                    
                    # Upload file to Gemini
                    uploaded_file = genai.upload_file(tmp_file_path, display_name=file_info['filename'])
                    content_parts.append(uploaded_file)
                    
                    # Clean up temporary file
                    os.unlink(tmp_file_path)
                    
                    print(f"   📎 Uploaded {file_info['filename']} to Gemini")
                    
                except Exception as e:
                    print(f"   ⚠️  Failed to upload {file_info.get('filename', 'unknown file')} to Gemini: {e}")
                    # Continue without this file
                    pass
        
        response = self.model.generate_content(content_parts, stream=False)
        return response.text

    def generate_content_stream(self, prompt: str, files: Optional[List[Dict[str, Any]]] = None) -> Generator[Any, None, None]:
        """Generates a streaming response from Gemini with optional file attachments."""
        content_parts = [prompt]
        
        # Upload files to Gemini if provided
        if files:
            for file_info in files:
                try:
                    # Write file content to temporary file
                    with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_info['extension'].lstrip('.')}") as tmp_file:
                        tmp_file.write(file_info['content'])
                        tmp_file_path = tmp_file.name
                    
                    # Upload file to Gemini
                    uploaded_file = genai.upload_file(tmp_file_path, display_name=file_info['filename'])
                    content_parts.append(uploaded_file)
                    
                    # Clean up temporary file
                    os.unlink(tmp_file_path)
                    
                    print(f"   📎 Uploaded {file_info['filename']} to Gemini")
                    
                except Exception as e:
                    print(f"   ⚠️  Failed to upload {file_info.get('filename', 'unknown file')} to Gemini: {e}")
                    # Continue without this file
                    pass
        
        return self.model.generate_content(content_parts, stream=True)
    
    def _generate_content_impl(self, prompt: str) -> str:
        """Fallback implementation for clients that don't support files."""
        response = self.model.generate_content(prompt, stream=False)
        return response.text

    def _generate_content_stream_impl(self, prompt: str) -> Generator[Any, None, None]:
        """Fallback implementation for clients that don't support files."""
        return self.model.generate_content(prompt, stream=True)