import os
import base64
import json
import logging
from typing import Dict, List, Optional, Tuple
import google.generativeai as genai
from PIL import Image
import io
import uuid
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('image_enhancer.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class ImageEnhancer:
    """
    AI Image Enhancement using Google Gemini 2.0 Flash Preview Image Generation
    Generates 5 enhanced versions of uploaded images
    """
    
    def __init__(self):
        """Initialize the image enhancer with Gemini"""
        logger.info("Initializing ImageEnhancer with Gemini...")
        
        api_key = os.getenv('GEMINI_API_KEY')
        if not api_key:
            logger.error("GEMINI_API_KEY not found in environment variables")
            raise ValueError("GEMINI_API_KEY not found in environment variables")
        
        logger.info("Gemini API key found, initializing client...")
        genai.configure(api_key=api_key)
        
        # Try different models for image generation
        try:
            # First try the experimental model
            self.model = genai.GenerativeModel('gemini-2.0-flash-exp')
            logger.info("Using gemini-2.0-flash-exp model")
        except Exception as e:
            logger.warning(f"Failed to initialize gemini-2.0-flash-exp: {e}")
            try:
                # Fallback to standard model
                self.model = genai.GenerativeModel('gemini-1.5-flash')
                logger.info("Using gemini-1.5-flash model")
            except Exception as e2:
                logger.error(f"Failed to initialize any Gemini model: {e2}")
                raise e2
        
        logger.info("ImageEnhancer initialization complete")
    
    def enhance_image(self, image_path: str, user_preferences: Optional[Dict] = None) -> Dict:
        """
        Generate 5 enhanced versions of an image
        
        Args:
            image_path: Path to the original image
            user_preferences: User preferences for enhancement style
            
        Returns:
            Dict containing enhanced images and metadata
        """
        logger.info(f"Starting image enhancement for: {image_path}")
        
        try:
            # Load original image
            original_image = Image.open(image_path)
            logger.info(f"Original image loaded. Size: {original_image.size}, Mode: {original_image.mode}")
            
            # Analyze original image to understand what needs enhancement
            analysis = self._analyze_image_for_enhancement(original_image)
            logger.info(f"Image analysis complete: {analysis['issues_found']}")
            
            # Generate enhancement prompts based on analysis and user preferences
            enhancement_prompts = self._generate_enhancement_prompts(analysis, user_preferences)
            logger.info(f"Generated {len(enhancement_prompts)} enhancement prompts")
            
            # Generate enhanced images
            enhanced_images = []
            for i, prompt in enumerate(enhancement_prompts):
                logger.info(f"Generating enhanced image {i+1}/5...")
                enhanced_image = self._generate_enhanced_image(original_image, prompt, i+1)
                if enhanced_image:
                    enhanced_images.append(enhanced_image)
            
            logger.info(f"Successfully generated {len(enhanced_images)} enhanced images")
            
            return {
                'status': 'success',
                'original_image': {
                    'path': image_path,
                    'size': original_image.size,
                    'analysis': analysis
                },
                'enhanced_images': enhanced_images,
                'total_generated': len(enhanced_images),
                'enhancement_prompts': enhancement_prompts
            }
            
        except Exception as e:
            logger.error(f"Image enhancement failed: {str(e)}", exc_info=True)
            return {
                'status': 'error',
                'error': f'Enhancement failed: {str(e)}',
                'enhanced_images': []
            }
    
    def _analyze_image_for_enhancement(self, image: Image) -> Dict:
        """
        Analyze image to identify enhancement opportunities
        """
        logger.info("Analyzing image for enhancement opportunities...")
        
        prompt = """
        Analyze this image and identify specific issues that could be improved for social media.
        Focus on technical and compositional issues that can be enhanced without changing personal appearance.
        
        Look for:
        1. Blurry or out-of-focus areas
        2. Closed eyes in group photos
        3. Unwanted objects (fingers, passing people, etc.)
        4. Poor lighting or exposure
        5. Composition issues
        6. Color balance problems
        7. Noise or grain
        8. Cropping opportunities
        
        Return ONLY a JSON object with:
        {
            "issues_found": ["list of specific issues"],
            "enhancement_priorities": ["ordered list of what to fix first"],
            "overall_quality": "good/medium/poor",
            "main_subject": "description of main subject",
            "background": "description of background",
            "lighting": "description of lighting conditions"
        }
        """
        
        try:
            response = self._get_gemini_response(image, prompt)
            analysis = json.loads(self._clean_json_response(response))
            logger.info(f"Image analysis: {analysis['issues_found']} issues found")
            return analysis
        except Exception as e:
            logger.error(f"Image analysis failed: {e}")
            return {
                "issues_found": ["general enhancement needed"],
                "enhancement_priorities": ["improve overall quality"],
                "overall_quality": "medium",
                "main_subject": "person or object",
                "background": "various",
                "lighting": "mixed"
            }
    
    def _generate_enhancement_prompts(self, analysis: Dict, user_preferences: Optional[Dict] = None) -> List[str]:
        """
        Generate 5 different enhancement prompts focused on fixing imperfections without changing personal appearance
        """
        logger.info("Generating enhancement prompts...")
        
        issues = analysis.get('issues_found', [])
        priorities = analysis.get('enhancement_priorities', [])
        
        # Base enhancement focus areas (as per transcript requirements)
        enhancement_focuses = [
            "fix blurry areas and improve overall sharpness and focus",
            "correct closed eyes and improve facial clarity and expressions", 
            "remove unwanted objects and clean up the background",
            "enhance lighting, exposure, and color balance",
            "improve composition, framing, and overall image quality"
        ]
        
        prompts = []
        for i, focus in enumerate(enhancement_focuses):
            # Create specific enhancement prompt focused on fixing imperfections
            prompt = f"Enhance this image by {focus}. "
            
            # Add specific fixes based on analysis
            if issues:
                specific_fixes = []
                for issue in issues[:3]:  # Focus on top 3 issues
                    if 'blur' in issue.lower() or 'focus' in issue.lower():
                        specific_fixes.append("fix any blurry or out-of-focus areas")
                    elif 'eye' in issue.lower():
                        specific_fixes.append("ensure all eyes are open and clear")
                    elif 'light' in issue.lower() or 'exposure' in issue.lower():
                        specific_fixes.append("improve lighting and exposure")
                    elif 'color' in issue.lower():
                        specific_fixes.append("enhance color balance and vibrancy")
                    elif 'noise' in issue.lower():
                        specific_fixes.append("reduce noise and improve clarity")
                    elif 'object' in issue.lower() or 'finger' in issue.lower():
                        specific_fixes.append("remove unwanted objects or distractions")
                
                if specific_fixes:
                    prompt += "Specifically address: " + ", ".join(specific_fixes) + ". "
            
            # Add user preference context for styling (not personal appearance)
            if user_preferences:
                aesthetic = user_preferences.get('aesthetic', '')
                niche = user_preferences.get('niche', '')
                if aesthetic and niche:
                    prompt += f"Apply {aesthetic} aesthetic styling for {niche} content. "
            
            # Critical instructions: NO personal appearance changes
            prompt += "IMPORTANT: Do NOT alter personal appearance, body shape, facial features, or make anyone look like a celebrity. Only fix technical image issues like blur, lighting, composition, and remove unwanted objects. Maintain the original person's appearance exactly as they are."
            
            prompts.append(prompt)
        
        logger.info(f"Generated {len(prompts)} enhancement prompts focused on fixing imperfections")
        return prompts
    
    def _generate_enhanced_image(self, original_image: Image, prompt: str, version: int) -> Optional[Dict]:
        """
        Generate a single enhanced image using Gemini 2.0 Flash Preview Image Generation
        """
        try:
            logger.info(f"Generating enhanced image version {version} with Gemini...")
            
            # Prepare the image for Gemini
            img_buffer = io.BytesIO()
            original_image.save(img_buffer, format='JPEG', quality=95)
            img_buffer.seek(0)
            
            # Create the content for Gemini (image + text prompt)
            content = [
                {
                    "mime_type": "image/jpeg",
                    "data": img_buffer.getvalue()
                },
                prompt
            ]
            
            # Generate enhanced image using Gemini 2.0 Flash Preview
            logger.info(f"Sending enhancement request to Gemini for version {version}...")
            
            # Try to generate content with image generation capability
            try:
                response = self.model.generate_content(
                    content,
                    generation_config=genai.types.GenerationConfig(
                        temperature=0.7,
                        max_output_tokens=8192,
                    )
                )
            except Exception as e:
                logger.error(f"Gemini generation failed: {e}")
                # Fallback to placeholder enhancement
                return self._create_enhanced_placeholder(original_image, prompt, version)
            
            # Log the response structure for debugging
            logger.info(f"Response type: {type(response)}")
            logger.info(f"Response parts: {len(response.parts) if hasattr(response, 'parts') else 'No parts'}")
            
            # Check if response contains an image
            if hasattr(response, 'parts') and response.parts:
                for i, part in enumerate(response.parts):
                    logger.info(f"Part {i}: {type(part)}")
                    logger.info(f"Part {i} attributes: {dir(part)}")
                    
                    # Check for inline_data (image data)
                    if hasattr(part, 'inline_data') and part.inline_data:
                        logger.info(f"Found inline_data in part {i}")
                        enhanced_image_path = self._save_gemini_image(part.inline_data.data, version)
                        
                        if enhanced_image_path:
                            logger.info(f"Enhanced image {version} saved to: {enhanced_image_path}")
                            return {
                                'version': version,
                                'prompt': prompt,
                                'image_path': enhanced_image_path,
                                'generation_method': 'gemini-2.0-flash-preview'
                            }
                        else:
                            logger.error(f"Failed to save enhanced image {version}")
                            return None
                    
                    # Check for text content (might contain image generation instructions)
                    elif hasattr(part, 'text') and part.text:
                        logger.info(f"Part {i} contains text: {part.text[:200]}...")
                        
                        # If Gemini returns text instead of image, try to extract image generation instructions
                        if "generate" in part.text.lower() or "create" in part.text.lower():
                            logger.info("Gemini returned text instructions instead of image")
                            # Use the text as enhancement instructions
                            return self._create_enhanced_placeholder(original_image, prompt, version, part.text)
            
            # If no image was generated, log the full response for debugging
            logger.error(f"No image generated in response for version {version}")
            logger.error(f"Full response: {response}")
            
            # Check if Gemini provided enhancement instructions in text
            enhancement_instructions = None
            if hasattr(response, 'parts') and response.parts:
                for part in response.parts:
                    if hasattr(part, 'text') and part.text:
                        enhancement_instructions = part.text
                        logger.info(f"Gemini provided enhancement instructions: {enhancement_instructions[:200]}...")
                        break
            
            # Try alternative approach - create enhanced image based on instructions
            return self._create_enhanced_placeholder(original_image, prompt, version, enhancement_instructions)
                
        except Exception as e:
            logger.error(f"Failed to generate enhanced image {version}: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def _save_gemini_image(self, image_data: bytes, version: int) -> Optional[str]:
        """
        Save Gemini generated image to local storage
        """
        try:
            # Create enhanced images directory
            enhanced_dir = "enhanced_images"
            os.makedirs(enhanced_dir, exist_ok=True)
            
            # Generate unique filename
            filename = f"enhanced_v{version}_{uuid.uuid4().hex[:8]}.jpg"
            filepath = os.path.join(enhanced_dir, filename)
            
            # Save image data directly
            with open(filepath, 'wb') as f:
                f.write(image_data)
            
            logger.info(f"Gemini enhanced image {version} saved to: {filepath}")
            return filepath
            
        except Exception as e:
            logger.error(f"Failed to save Gemini image: {e}")
            return None
    
    def _create_enhanced_placeholder(self, original_image: Image, prompt: str, version: int, instructions: str = None) -> Optional[Dict]:
        """
        Create a placeholder enhanced image when Gemini doesn't generate images
        This applies basic image processing to simulate enhancement
        """
        try:
            logger.info(f"Creating enhanced placeholder for version {version}...")
            
            # Create enhanced images directory
            enhanced_dir = "enhanced_images"
            os.makedirs(enhanced_dir, exist_ok=True)
            
            # Generate unique filename
            filename = f"enhanced_v{version}_{uuid.uuid4().hex[:8]}.jpg"
            filepath = os.path.join(enhanced_dir, filename)
            
            # Apply basic enhancements to simulate AI enhancement
            enhanced_image = original_image.copy()
            
            # Apply different enhancements based on version to create variety
            enhancement_factors = {
                1: {'brightness': 1.1, 'contrast': 1.05, 'color': 1.1, 'sharpness': True},
                2: {'brightness': 1.05, 'contrast': 1.1, 'color': 1.05, 'sharpness': True},
                3: {'brightness': 1.15, 'contrast': 1.0, 'color': 1.15, 'sharpness': False},
                4: {'brightness': 1.0, 'contrast': 1.15, 'color': 1.0, 'sharpness': True},
                5: {'brightness': 1.08, 'contrast': 1.08, 'color': 1.08, 'sharpness': True}
            }
            
            factors = enhancement_factors.get(version, enhancement_factors[1])
            
            # Apply enhancements
            from PIL import ImageEnhance, ImageFilter
            
            # Brightness
            if factors['brightness'] != 1.0:
                enhancer = ImageEnhance.Brightness(enhanced_image)
                enhanced_image = enhancer.enhance(factors['brightness'])
            
            # Contrast
            if factors['contrast'] != 1.0:
                enhancer = ImageEnhance.Contrast(enhanced_image)
                enhanced_image = enhancer.enhance(factors['contrast'])
            
            # Color
            if factors['color'] != 1.0:
                enhancer = ImageEnhance.Color(enhanced_image)
                enhanced_image = enhancer.enhance(factors['color'])
            
            # Sharpness
            if factors['sharpness']:
                enhanced_image = enhanced_image.filter(ImageFilter.SHARPEN)
            
            # Apply specific enhancements based on prompt
            if "blur" in prompt.lower() or "sharp" in prompt.lower():
                enhanced_image = enhanced_image.filter(ImageFilter.SHARPEN)
            
            if "light" in prompt.lower() or "exposure" in prompt.lower():
                enhancer = ImageEnhance.Brightness(enhanced_image)
                enhanced_image = enhancer.enhance(1.1)
            
            if "color" in prompt.lower():
                enhancer = ImageEnhance.Color(enhanced_image)
                enhanced_image = enhancer.enhance(1.1)
            
            # Save the enhanced image
            enhanced_image.save(filepath, 'JPEG', quality=95)
            
            logger.info(f"Enhanced placeholder {version} saved to: {filepath}")
            return {
                'version': version,
                'prompt': prompt,
                'image_path': filepath,
                'generation_method': 'ai-enhanced-placeholder',
                'enhancement_factors': factors,
                'gemini_instructions': instructions[:200] if instructions else None
            }
            
        except Exception as e:
            logger.error(f"Failed to create enhanced placeholder {version}: {e}")
            return None
    
    def _get_gemini_response(self, image: Image, prompt: str) -> str:
        """Get response from Gemini for image analysis"""
        try:
            # Prepare the image for Gemini
            img_buffer = io.BytesIO()
            image.save(img_buffer, format='JPEG')
            img_buffer.seek(0)
            
            # Create the content for Gemini (image + text prompt)
            content = [
                {
                    "mime_type": "image/jpeg",
                    "data": img_buffer.getvalue()
                },
                prompt
            ]
            
            # Get response from Gemini
            response = self.model.generate_content(
                content,
                generation_config=genai.types.GenerationConfig(
                    temperature=0.3,
                    max_output_tokens=1000,
                )
            )
            
            return response.text
            
        except Exception as e:
            logger.error(f"Gemini API error: {e}")
            return "{}"
    
    def _clean_json_response(self, response: str) -> str:
        """Clean JSON response by removing markdown formatting"""
        cleaned_response = response.strip()
        if cleaned_response.startswith('```json'):
            cleaned_response = cleaned_response[7:]
        if cleaned_response.startswith('```'):
            cleaned_response = cleaned_response[3:]
        if cleaned_response.endswith('```'):
            cleaned_response = cleaned_response[:-3]
        
        return cleaned_response.strip()

