#!/usr/bin/env python3
"""
Quick test to verify the image enhancement fix works
"""

import os
import sys
import logging
from image_enhancer import ImageEnhancer
from PIL import Image

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_enhancement_fix():
    """Test the fixed image enhancement"""
    print("🧪 Testing Image Enhancement Fix")
    print("=" * 40)
    
    # Check Gemini API key
    api_key = os.getenv('GEMINI_API_KEY')
    if not api_key:
        print("❌ GEMINI_API_KEY not found")
        print("   The system will use placeholder enhancement")
    else:
        print("✅ Gemini API key found")
    
    # Initialize enhancer
    try:
        print("1. Initializing Image Enhancer...")
        enhancer = ImageEnhancer()
        print("   ✅ Enhancer initialized successfully")
    except Exception as e:
        print(f"   ❌ Failed to initialize enhancer: {e}")
        return False
    
    # Create a simple test image
    print("\n2. Creating test image...")
    test_image = Image.new('RGB', (100, 100), color='lightblue')
    test_image_path = "test_fix_image.jpg"
    test_image.save(test_image_path, 'JPEG')
    print(f"   ✅ Test image created: {test_image_path}")
    
    # Test enhancement
    print("\n3. Testing enhancement...")
    try:
        result = enhancer.enhance_image(test_image_path)
        
        if result['status'] == 'success':
            print(f"   ✅ Enhancement completed!")
            print(f"   📊 Generated {result['total_generated']} enhanced images")
            
            for i, enhanced_img in enumerate(result['enhanced_images'], 1):
                method = enhanced_img.get('generation_method', 'unknown')
                print(f"     Version {i}: {method}")
                
                if method == 'ai-enhanced-placeholder':
                    factors = enhanced_img.get('enhancement_factors', {})
                    print(f"       Enhancement factors: {factors}")
            
            return True
        else:
            print(f"   ❌ Enhancement failed: {result.get('error', 'Unknown error')}")
            return False
            
    except Exception as e:
        print(f"   ❌ Test failed: {e}")
        return False
    
    finally:
        # Clean up
        if os.path.exists(test_image_path):
            os.remove(test_image_path)
            print(f"\n   🧹 Cleaned up test image")

def main():
    """Main test function"""
    print("🚀 Viral Velocity - Enhancement Fix Test")
    print("=" * 50)
    
    success = test_enhancement_fix()
    
    print("\n" + "=" * 50)
    if success:
        print("✅ Enhancement fix test completed successfully!")
        print("🎉 The system now generates enhanced images")
        print("📝 Features:")
        print("   - Gemini integration with fallback")
        print("   - AI-enhanced placeholder images")
        print("   - 5 different enhancement variations")
        print("   - Proper error handling")
    else:
        print("❌ Enhancement fix test failed")
        print("📝 Check the logs above for details")

if __name__ == "__main__":
    main()