#!/usr/bin/env python3
"""
Image Understanding + Modification + Verification System

This module combines:
1. Google Gemini AI - Image understanding and analysis
2. FAL AI - Image modification using multiple models
3. Verification - Compare before/after analysis

Workflow: Understand → Modify → Verify

Features:
- Multi-model image modification (Photon, Kontext, SeedEdit)
- AI-powered image analysis and understanding
- Before/after comparison and verification
- Intelligent modification suggestions
- Comprehensive error handling and validation

Author: AI Assistant
Date: 2024
"""

import os
import sys
from pathlib import Path
from typing import Dict, Any, Optional, List, Tuple
import json
import shutil
from datetime import datetime

# Add paths for both systems
current_dir = Path(__file__).parent
video_tools_dir = current_dir.parent
fal_image_dir = video_tools_dir.parent / "fal_image_to_image"
sys.path.insert(0, str(video_tools_dir))
sys.path.insert(0, str(fal_image_dir))

# Import video_tools components (Gemini understanding)
from video_utils.video_understanding import GeminiVideoAnalyzer, save_analysis_result

# Import fal_image_to_image components (FAL modification)
try:
    # Try to import from the fal_image_to_image directory
    from fal_image_to_image_generator import FALImageToImageGenerator
    FAL_AVAILABLE = True
except ImportError:
    FAL_AVAILABLE = False
    print("⚠️  FAL Image-to-Image module not available")

class ImageModifyVerifySystem:
    """
    Comprehensive system for image understanding, modification, and verification.
    
    Combines Google Gemini AI for understanding with FAL AI for modification.
    """
    
    def __init__(self, output_dir: Optional[Path] = None):
        """
        Initialize the image modification system.
        
        Args:
            output_dir: Directory for saving results (default: current/output)
        """
        self.output_dir = output_dir or (Path(__file__).parent.parent / "output")
        self.output_dir.mkdir(exist_ok=True)
        
        # Initialize analyzers
        try:
            self.gemini_analyzer = GeminiVideoAnalyzer()
            self.gemini_available = True
        except Exception as e:
            print(f"⚠️  Gemini analyzer not available: {e}")
            self.gemini_available = False
        
        if FAL_AVAILABLE:
            try:
                self.fal_generator = FALImageToImageGenerator()
                self.fal_available = True
            except Exception as e:
                print(f"⚠️  FAL generator not available: {e}")
                self.fal_available = False
        else:
            self.fal_available = False
    
    def check_requirements(self) -> Tuple[bool, str]:
        """Check if both systems are available."""
        if not self.gemini_available:
            return False, "Google Gemini API not configured"
        if not self.fal_available:
            return False, "FAL AI not configured"
        return True, "All systems ready"
    
    def understand_image(self, image_path: Path, analysis_types: List[str] = None) -> Dict[str, Any]:
        """
        Analyze and understand the image using Google Gemini.
        
        Args:
            image_path: Path to image file
            analysis_types: List of analysis types ['description', 'objects', 'composition', 'text']
        
        Returns:
            Comprehensive analysis results
        """
        if not self.gemini_available:
            raise RuntimeError("Gemini analyzer not available")
        
        if analysis_types is None:
            analysis_types = ['description', 'objects', 'composition']
        
        print(f"🔍 Analyzing image: {image_path.name}")
        
        results = {
            'file_path': str(image_path),
            'timestamp': datetime.now().isoformat(),
            'analysis': {}
        }
        
        for analysis_type in analysis_types:
            print(f"   📋 Running {analysis_type} analysis...")
            try:
                if analysis_type == 'description':
                    result = self.gemini_analyzer.describe_image(image_path, detailed=True)
                    results['analysis']['description'] = result['description']
                
                elif analysis_type == 'objects':
                    result = self.gemini_analyzer.detect_objects(image_path, detailed=True)
                    results['analysis']['objects'] = result['objects']
                
                elif analysis_type == 'composition':
                    result = self.gemini_analyzer.analyze_composition(image_path, detailed=True)
                    results['analysis']['composition'] = result['composition_analysis']
                
                elif analysis_type == 'text':
                    result = self.gemini_analyzer.extract_text_from_image(image_path)
                    results['analysis']['text'] = result['extracted_text']
                
                elif analysis_type == 'classification':
                    result = self.gemini_analyzer.classify_image(image_path)
                    results['analysis']['classification'] = result['classification']
                
            except Exception as e:
                print(f"   ❌ {analysis_type} analysis failed: {e}")
                results['analysis'][analysis_type] = f"Error: {e}"
        
        return results
    
    def suggest_modifications(self, understanding: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Generate intelligent modification suggestions based on image understanding.
        
        Args:
            understanding: Results from understand_image()
        
        Returns:
            List of modification suggestions
        """
        suggestions = []
        
        # Extract key information
        description = understanding['analysis'].get('description', '')
        objects = understanding['analysis'].get('objects', '')
        composition = understanding['analysis'].get('composition', '')
        
        # Analyze content and suggest modifications
        description_lower = description.lower()
        objects_lower = objects.lower()
        
        # Lighting suggestions
        if any(word in description_lower for word in ['dark', 'dim', 'shadow', 'low light']):
            suggestions.append({
                'type': 'lighting',
                'prompt': 'Brighten the image with natural lighting',
                'model': 'photon',
                'strength': 0.6,
                'reason': 'Image appears to have low lighting'
            })
        
        # Color enhancement suggestions
        if any(word in description_lower for word in ['dull', 'muted', 'faded', 'pale']):
            suggestions.append({
                'type': 'color_enhancement',
                'prompt': 'Enhance colors and vibrancy, make more vivid and saturated',
                'model': 'photon',
                'strength': 0.5,
                'reason': 'Colors could be enhanced'
            })
        
        # Background suggestions
        if 'cluttered' in description_lower or 'messy' in description_lower:
            suggestions.append({
                'type': 'background_cleanup',
                'prompt': 'Clean up the background, remove clutter and distractions',
                'model': 'kontext',
                'inference_steps': 30,
                'reason': 'Background appears cluttered'
            })
        
        # Style suggestions
        if any(word in description_lower for word in ['portrait', 'person', 'face']):
            suggestions.append({
                'type': 'portrait_enhancement',
                'prompt': 'Enhance portrait quality, improve skin tone and facial features',
                'model': 'photon',
                'strength': 0.4,
                'reason': 'Portrait detected'
            })
        
        # Scene-specific suggestions
        if any(word in description_lower for word in ['landscape', 'outdoor', 'nature']):
            suggestions.append({
                'type': 'landscape_enhancement',
                'prompt': 'Enhance natural colors, improve sky and vegetation',
                'model': 'photon',
                'strength': 0.5,
                'reason': 'Landscape scene detected'
            })
        
        # Quality improvements
        if any(word in description_lower for word in ['blurry', 'noisy', 'grainy', 'low quality']):
            suggestions.append({
                'type': 'quality_improvement',
                'prompt': 'Improve image quality, reduce noise and enhance sharpness',
                'model': 'kontext',
                'inference_steps': 35,
                'reason': 'Image quality could be improved'
            })
        
        # Default suggestion if none match
        if not suggestions:
            suggestions.append({
                'type': 'general_enhancement',
                'prompt': 'Enhance overall image quality and visual appeal',
                'model': 'photon',
                'strength': 0.5,
                'reason': 'General enhancement'
            })
        
        return suggestions
    
    def modify_image(self, image_path: Path, modification: Dict[str, Any]) -> Optional[Path]:
        """
        Modify image using FAL AI based on modification parameters.
        
        Args:
            image_path: Path to original image
            modification: Modification parameters
        
        Returns:
            Path to modified image or None if failed
        """
        if not self.fal_available:
            raise RuntimeError("FAL generator not available")
        
        print(f"🎨 Modifying image with {modification['model']} model...")
        print(f"   📝 Prompt: {modification['prompt']}")
        
        try:
            # Upload image to get URL (FAL API requires image URL)
            print(f"   📤 Uploading image for modification...")
            with open(image_path, 'rb') as f:
                import base64
                image_data = base64.b64encode(f.read()).decode()
                image_url = f"data:image/jpeg;base64,{image_data}"
            
            # Prepare parameters based on model
            if modification['model'] in ['photon', 'photon_base']:
                result = self.fal_generator.modify_image(
                    image_url=image_url,
                    prompt=modification['prompt'],
                    model=modification['model'],
                    strength=modification.get('strength', 0.5),
                    aspect_ratio=modification.get('aspect_ratio', '1:1')
                )
            
            elif modification['model'] == 'kontext':
                result = self.fal_generator.modify_image(
                    image_url=image_url,
                    prompt=modification['prompt'],
                    model='kontext',
                    num_inference_steps=modification.get('inference_steps', 28),
                    guidance_scale=modification.get('guidance_scale', 2.5)
                )
            
            else:
                raise ValueError(f"Unsupported model: {modification['model']}")
            
            if result and 'image_url' in result:
                # Download modified image
                import requests
                response = requests.get(result['image_url'])
                if response.status_code == 200:
                    # Create unique filename
                    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                    modified_path = self.output_dir / f"{image_path.stem}_modified_{timestamp}.png"
                    
                    with open(modified_path, 'wb') as f:
                        f.write(response.content)
                    
                    print(f"✅ Modified image saved: {modified_path.name}")
                    return modified_path
            
            print("❌ No image URL in result")
            return None
            
        except Exception as e:
            print(f"❌ Modification failed: {e}")
            return None
    
    def verify_modification(self, original_path: Path, modified_path: Path) -> Dict[str, Any]:
        """
        Verify modification by comparing original and modified images.
        
        Args:
            original_path: Path to original image
            modified_path: Path to modified image
        
        Returns:
            Comparison analysis
        """
        print("🔍 Verifying modification...")
        
        # Analyze both images
        original_analysis = self.understand_image(original_path, ['description', 'objects'])
        modified_analysis = self.understand_image(modified_path, ['description', 'objects'])
        
        # Compare results
        verification = {
            'timestamp': datetime.now().isoformat(),
            'original': {
                'path': str(original_path),
                'analysis': original_analysis['analysis']
            },
            'modified': {
                'path': str(modified_path),
                'analysis': modified_analysis['analysis']
            },
            'comparison': {}
        }
        
        # Generate comparison insights using Gemini
        try:
            comparison_prompt = f"""Compare these two image descriptions and analyze the changes:

ORIGINAL IMAGE:
{original_analysis['analysis'].get('description', 'No description')}

MODIFIED IMAGE:  
{modified_analysis['analysis'].get('description', 'No description')}

Please analyze:
1. What specific changes were made?
2. How successful was the modification?
3. What improvements are visible?
4. What could be improved further?
5. Rate the modification success (1-10)

Provide a detailed comparison."""

            # Note: This would require a text-only Gemini call
            # For now, we'll do a simple text comparison
            verification['comparison']['summary'] = "Modification completed - manual review recommended"
            verification['comparison']['success_rating'] = "Pending manual review"
            
        except Exception as e:
            verification['comparison']['error'] = str(e)
        
        return verification
    
    def complete_workflow(self, image_path: Path, custom_prompt: Optional[str] = None, 
                         model: str = 'photon') -> Dict[str, Any]:
        """
        Complete workflow: Understand → Modify → Verify
        
        Args:
            image_path: Path to image file
            custom_prompt: Custom modification prompt (optional)
            model: Model to use for modification
        
        Returns:
            Complete workflow results
        """
        workflow_results = {
            'workflow_id': f"workflow_{datetime.now().strftime('%Y%m%d_%H%M%S')}",
            'original_image': str(image_path),
            'timestamp': datetime.now().isoformat(),
            'steps': {}
        }
        
        try:
            # Step 1: Understand
            print("🎯 STEP 1: Understanding Image")
            understanding = self.understand_image(image_path)
            workflow_results['steps']['understanding'] = understanding
            
            # Save understanding results
            understanding_file = self.output_dir / f"{image_path.stem}_understanding.json"
            save_analysis_result(understanding, understanding_file)
            
            # Step 2: Generate or use modification
            print("\n🎯 STEP 2: Planning Modification")
            if custom_prompt:
                modification = {
                    'type': 'custom',
                    'prompt': custom_prompt,
                    'model': model,
                    'strength': 0.5 if model in ['photon', 'photon_base'] else None,
                    'inference_steps': 28 if model == 'kontext' else None,
                    'reason': 'User provided custom prompt'
                }
            else:
                suggestions = self.suggest_modifications(understanding)
                modification = suggestions[0] if suggestions else None
            
            if not modification:
                raise ValueError("No modification could be planned")
            
            workflow_results['steps']['modification_plan'] = modification
            print(f"   📝 Planned: {modification['reason']}")
            print(f"   🎨 Prompt: {modification['prompt']}")
            
            # Step 3: Modify
            print("\n🎯 STEP 3: Modifying Image")
            modified_path = self.modify_image(image_path, modification)
            
            if not modified_path:
                raise RuntimeError("Image modification failed")
            
            workflow_results['steps']['modified_image'] = str(modified_path)
            
            # Step 4: Verify
            print("\n🎯 STEP 4: Verifying Results")
            verification = self.verify_modification(image_path, modified_path)
            workflow_results['steps']['verification'] = verification
            
            # Save complete results
            workflow_file = self.output_dir / f"{image_path.stem}_workflow_complete.json"
            save_analysis_result(workflow_results, workflow_file)
            
            print(f"\n🎉 Workflow Complete!")
            print(f"   📁 Original: {image_path.name}")
            print(f"   🎨 Modified: {modified_path.name}")
            print(f"   📊 Results: {workflow_file.name}")
            
            return workflow_results
            
        except Exception as e:
            workflow_results['error'] = str(e)
            print(f"\n❌ Workflow failed: {e}")
            return workflow_results

def main():
    """CLI interface for image modification workflow."""
    import argparse
    
    parser = argparse.ArgumentParser(description="Image Understanding + Modification + Verification")
    parser.add_argument("image", help="Path to image file")
    parser.add_argument("--prompt", help="Custom modification prompt")
    parser.add_argument("--model", choices=['photon', 'photon_base', 'kontext'], 
                       default='photon', help="Modification model")
    parser.add_argument("--output", help="Output directory")
    parser.add_argument("--understand-only", action="store_true", 
                       help="Only analyze image, don't modify")
    
    args = parser.parse_args()
    
    # Initialize system
    output_dir = Path(args.output) if args.output else None
    system = ImageModifyVerifySystem(output_dir)
    
    # Check requirements
    ready, message = system.check_requirements()
    if not ready:
        print(f"❌ System not ready: {message}")
        return 1
    
    image_path = Path(args.image)
    if not image_path.exists():
        print(f"❌ Image file not found: {image_path}")
        return 1
    
    if args.understand_only:
        # Only understand the image
        print("🔍 Understanding image only (no modification)")
        understanding = system.understand_image(image_path)
        
        print("\n📋 Image Analysis:")
        for analysis_type, result in understanding['analysis'].items():
            print(f"\n{analysis_type.upper()}:")
            preview = result[:200] + "..." if len(result) > 200 else result
            print(f"   {preview}")
    else:
        # Complete workflow
        result = system.complete_workflow(image_path, args.prompt, args.model)
        
        if 'error' not in result:
            print(f"\n✅ Success! Check output directory: {system.output_dir}")
        else:
            print(f"\n❌ Failed: {result['error']}")
            return 1
    
    return 0

if __name__ == "__main__":
    exit(main())