#!/usr/bin/env python3
"""
AI Background Generator Module - Gradient-Only Version
Handles background generation with smart gradient fallbacks when AI libraries conflict.
"""

import os
import sys
import tempfile
import random
import logging
from pathlib import Path
from typing import Optional, Tuple
import io
import base64

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

# Safe imports only
PIL_AVAILABLE = False
TORCH_AVAILABLE = False
DIFFUSERS_AVAILABLE = False

try:
    from PIL import Image, ImageDraw, ImageFilter
    PIL_AVAILABLE = True
    logger.info("✅ PIL imported successfully")
except ImportError as e:
    logger.error(f"❌ PIL import failed: {e}")

# Check PyTorch availability but don't import if problematic
try:
    import torch
    TORCH_AVAILABLE = True
    logger.info("✅ PyTorch available")
    
    # Test for the specific custom_op issue
    if hasattr(torch.library, 'custom_op'):
        logger.info("✅ PyTorch custom_op available")
    else:
        logger.warning("⚠️  PyTorch custom_op not available - diffusers will likely fail")
        
except ImportError:
    logger.warning("⚠️  PyTorch not available")
except AttributeError:
    logger.warning("⚠️  PyTorch version incompatible (missing torch.library)")
except Exception as e:
    logger.warning(f"⚠️  PyTorch check failed: {e}")

# NEVER attempt to import diffusers if we detect the custom_op issue
FORCE_GRADIENT_MODE = False

if TORCH_AVAILABLE:
    try:
        # First, check if torch.library.custom_op exists
        import torch
        if not hasattr(torch.library, 'custom_op'):
            logger.warning("🔄 Force enabling gradient-only mode due to torch.library.custom_op missing")
            FORCE_GRADIENT_MODE = True
        else:
            # Only try diffusers if custom_op exists
            try:
                # Quick test import to see if diffusers will work
                import importlib.util
                spec = importlib.util.find_spec("diffusers")
                if spec is not None:
                    # Try a minimal import test
                    from diffusers import __version__
                    logger.info(f"✅ Diffusers {__version__} detected and compatible")
                    DIFFUSERS_AVAILABLE = True
                else:
                    logger.info("ℹ️  Diffusers not installed")
            except Exception as e:
                if "custom_op" in str(e):
                    logger.warning("🔄 Detected custom_op compatibility issue - using gradient-only mode")
                    FORCE_GRADIENT_MODE = True
                else:
                    logger.warning(f"🔄 Diffusers import issue: {e}")
    except Exception as e:
        logger.warning(f"🔄 PyTorch/Diffusers compatibility check failed: {e}")
        FORCE_GRADIENT_MODE = True
else:
    logger.info("ℹ️  Skipping diffusers check - PyTorch not available")

# Override diffusers availability if we're forcing gradient mode
if FORCE_GRADIENT_MODE:
    DIFFUSERS_AVAILABLE = False

class AIBackgroundGenerator:
    """
    AI Background Generator with intelligent gradient fallbacks
    """
    
    def __init__(self):
        self.pipeline = None
        self.device = "cpu"
        
        # Comprehensive color themes for gradient fallbacks
        self.color_themes = {
            # Blues
            'blue': [(64, 128, 255), (0, 64, 128)],
            'ocean': [(0, 119, 190), (0, 64, 128)], 
            'sky': [(135, 206, 250), (25, 25, 112)],
            'water': [(0, 191, 255), (0, 100, 200)],
            'azure': [(240, 255, 255), (0, 127, 255)],
            
            # Greens
            'green': [(34, 139, 34), (0, 100, 0)],
            'nature': [(107, 142, 35), (34, 139, 34)],
            'forest': [(34, 139, 34), (0, 50, 0)],
            'grass': [(124, 252, 0), (34, 139, 34)],
            'mint': [(152, 251, 152), (0, 128, 0)],
            
            # Professional/Business
            'professional': [(105, 105, 105), (169, 169, 169)],
            'office': [(192, 192, 192), (105, 105, 105)],
            'corporate': [(70, 130, 180), (25, 25, 112)],
            'business': [(47, 79, 79), (112, 128, 144)],
            'modern': [(95, 158, 160), (47, 79, 79)],
            
            # Dark themes
            'dark': [(64, 64, 64), (0, 0, 0)],
            'night': [(25, 25, 112), (0, 0, 0)],
            'black': [(105, 105, 105), (0, 0, 0)],
            'shadow': [(85, 85, 85), (0, 0, 0)],
            
            # Warm colors
            'warm': [(255, 140, 0), (255, 69, 0)],
            'sunset': [(255, 94, 77), (255, 154, 0)],
            'orange': [(255, 165, 0), (255, 69, 0)],
            'fire': [(255, 69, 0), (139, 0, 0)],
            
            # Cool colors
            'purple': [(147, 112, 219), (75, 0, 130)],
            'violet': [(138, 43, 226), (75, 0, 130)],
            'lavender': [(230, 230, 250), (147, 112, 219)],
            
            # Others
            'red': [(220, 20, 60), (139, 0, 0)],
            'pink': [(255, 182, 193), (255, 20, 147)],
            'yellow': [(255, 255, 0), (255, 215, 0)],
            'gold': [(255, 215, 0), (184, 134, 11)],
            
            # Technology/Digital
            'tech': [(0, 255, 255), (0, 0, 139)],
            'digital': [(138, 43, 226), (25, 25, 112)],
            'cyber': [(0, 255, 127), (0, 100, 0)],
            'neon': [(255, 20, 147), (138, 43, 226)],
            
            # Default
            'default': [(100, 149, 237), (65, 105, 225)]
        }
        
        # Only try to initialize diffusers if it's actually available
        if DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE:
            logger.info("🎨 Attempting to initialize AI pipeline...")
            self._init_diffusers()
        else:
            logger.info("🎨 Using gradient-only mode")
    
    def _init_diffusers(self):
        """Initialize the Stable Diffusion pipeline (only if safe)"""
        # This will only be called if we've verified diffusers works
        try:
            from diffusers import StableDiffusionPipeline
            
            model_id = "runwayml/stable-diffusion-v1-5"
            
            if torch.cuda.is_available():
                self.device = "cuda"
                logger.info("🚀 Using CUDA device")
            else:
                self.device = "cpu"
                logger.info("🖥️  Using CPU device")
            
            self.pipeline = StableDiffusionPipeline.from_pretrained(
                model_id,
                torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
                safety_checker=None,
                requires_safety_checker=False
            ).to(self.device)
            
            # Optimize for memory
            if self.device == "cuda":
                self.pipeline.enable_memory_efficient_attention()
                self.pipeline.enable_attention_slicing()
            
            logger.info("✅ AI pipeline initialized successfully")
            
        except Exception as e:
            logger.error(f"❌ Failed to initialize AI pipeline: {e}")
            self.pipeline = None
            global DIFFUSERS_AVAILABLE
            DIFFUSERS_AVAILABLE = False
    
    def _analyze_prompt_theme(self, prompt: str) -> str:
        """Analyze prompt to determine appropriate color theme"""
        prompt_lower = prompt.lower()
        
        # Direct theme matches first
        for theme in self.color_themes:
            if theme in prompt_lower:
                return theme
        
        # Keyword analysis
        keyword_map = {
            # Water/Ocean
            ('water', 'sea', 'lake', 'river', 'stream'): 'ocean',
            ('sky', 'cloud', 'air'): 'sky',
            
            # Nature
            ('tree', 'plant', 'garden', 'leaf', 'flower'): 'nature',
            ('grass', 'field', 'meadow'): 'grass',
            
            # Business
            ('business', 'meeting', 'work', 'conference'): 'business',
            ('company', 'enterprise', 'corporate'): 'corporate',
            
            # Time/Lighting
            ('evening', 'midnight', 'shadow', 'darkness'): 'night',
            ('fire', 'flame', 'autumn', 'fall'): 'fire',
            ('morning', 'sunrise', 'dawn'): 'warm',
            
            # Technology
            ('technology', 'computer', 'digital', 'software'): 'tech',
            ('cyber', 'virtual', 'matrix'): 'cyber',
            ('neon', 'electric', 'bright'): 'neon',
            
            # Emotions/Moods
            ('calm', 'peaceful', 'serene'): 'azure',
            ('energetic', 'vibrant', 'active'): 'orange',
            ('elegant', 'sophisticated', 'luxury'): 'purple',
            ('fresh', 'clean', 'pure'): 'mint',
        }
        
        for keywords, theme in keyword_map.items():
            if any(keyword in prompt_lower for keyword in keywords):
                return theme
        
        return 'default'
    
    def _create_gradient_background(self, width: int = 1024, height: int = 768, 
                                  theme: str = 'default') -> Image.Image:
        """Create a sophisticated gradient background"""
        if not PIL_AVAILABLE:
            raise RuntimeError("PIL is not available for gradient generation")
        
        # Get colors for the theme
        colors = self.color_themes.get(theme, self.color_themes['default'])
        color1, color2 = colors
        
        # Create base gradient
        image = Image.new('RGB', (width, height))
        draw = ImageDraw.Draw(image)
        
        # Create smooth vertical gradient
        for y in range(height):
            # Use smooth interpolation
            factor = y / height
            # Apply easing function for smoother gradients
            factor = factor * factor * (3.0 - 2.0 * factor)  # Smoothstep
            
            # Interpolate between colors
            r = int(color1[0] * (1 - factor) + color2[0] * factor)
            g = int(color1[1] * (1 - factor) + color2[1] * factor)
            b = int(color1[2] * (1 - factor) + color2[2] * factor)
            
            # Draw horizontal line
            draw.line([(0, y), (width, y)], fill=(r, g, b))
        
        # Add subtle texture and depth
        self._add_texture(image, theme)
        
        return image
    
    def _add_texture(self, image: Image.Image, theme: str):
        """Add subtle texture to make gradients more interesting"""
        width, height = image.size
        
        # Create texture overlay
        texture = Image.new('RGBA', (width, height), (0, 0, 0, 0))
        texture_draw = ImageDraw.Draw(texture)
        
        # Add different textures based on theme
        if theme in ['tech', 'digital', 'cyber']:
            # Add subtle grid pattern
            grid_size = 50
            for x in range(0, width, grid_size):
                texture_draw.line([(x, 0), (x, height)], fill=(255, 255, 255, 5))
            for y in range(0, height, grid_size):
                texture_draw.line([(0, y), (width, y)], fill=(255, 255, 255, 5))
        
        elif theme in ['nature', 'forest', 'grass']:
            # Add organic noise
            for _ in range(width * height // 200):
                x = random.randint(0, width - 1)
                y = random.randint(0, height - 1)
                size = random.randint(1, 3)
                brightness = random.randint(10, 30)
                texture_draw.ellipse([(x, y), (x+size, y+size)], 
                                   fill=(brightness, brightness, brightness, 20))
        
        else:
            # Add subtle noise for general texture
            for _ in range(width * height // 300):
                x = random.randint(0, width - 1)
                y = random.randint(0, height - 1)
                brightness = random.randint(-15, 15)
                alpha = random.randint(5, 15)
                texture_draw.point((x, y), fill=(brightness, brightness, brightness, alpha))
        
        # Apply texture
        image.paste(texture, (0, 0), texture)
        
        # Final smooth blur for professional look
        image = image.filter(ImageFilter.GaussianBlur(radius=0.8))
    
    def generate_background(self, prompt: str, width: int = 1024, height: int = 768,
                          guidance_scale: float = 7.5, num_inference_steps: int = 20) -> Optional[Image.Image]:
        """
        Generate a background image from a text prompt.
        Uses gradients with intelligent theming.
        """
        
        if not PIL_AVAILABLE:
            logger.error("❌ Cannot generate backgrounds - PIL not available")
            return None
        
        # Check if AI generation is possible and working
        if DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE and self.pipeline is not None:
            try:
                logger.info(f"🎨 Attempting AI generation for: '{prompt}'")
                
                enhanced_prompt = f"{prompt}, high quality, detailed, professional background, 8k"
                
                with torch.no_grad():
                    result = self.pipeline(
                        enhanced_prompt,
                        width=width,
                        height=height,
                        guidance_scale=guidance_scale,
                        num_inference_steps=num_inference_steps,
                        negative_prompt="low quality, blurry, distorted, watermark, text"
                    )
                
                if result.images and len(result.images) > 0:
                    logger.info("✅ AI background generated successfully")
                    return result.images[0]
                
            except Exception as e:
                logger.warning(f"⚠️  AI generation failed, using gradient: {e}")
        
        # Use gradient generation (main path for compatibility)
        logger.info(f"🎨 Creating gradient background for: '{prompt}'")
        theme = self._analyze_prompt_theme(prompt)
        logger.info(f"🎯 Selected theme: '{theme}'")
        
        try:
            image = self._create_gradient_background(width, height, theme)
            if image:
                logger.info("✅ Gradient background generated successfully")
                return image
            else:
                logger.error("❌ Gradient generation returned None")
                return None
        except Exception as e:
            logger.error(f"❌ Gradient generation failed: {e}")
            return None
    
    def save_background(self, image: Image.Image, output_path: str) -> bool:
        """Save the generated background image"""
        try:
            # Ensure directory exists
            Path(output_path).parent.mkdir(parents=True, exist_ok=True)
            
            # Save image with high quality
            image.save(output_path, format='PNG', quality=95, optimize=True)
            logger.info(f"💾 Background saved to: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"❌ Failed to save background: {e}")
            return False
    
    def get_background_base64(self, image: Image.Image) -> str:
        """Convert background image to base64 string"""
        try:
            buffer = io.BytesIO()
            image.save(buffer, format='PNG')
            img_str = base64.b64encode(buffer.getvalue()).decode()
            return img_str
        except Exception as e:
            logger.error(f"❌ Failed to convert to base64: {e}")
            return ""

# Convenience functions
def generate_ai_background(prompt: str, width: int = 1024, height: int = 768) -> Optional[Image.Image]:
    """
    Quick function to generate a background.
    
    Args:
        prompt: Text description of the desired background
        width: Image width in pixels
        height: Image height in pixels
    
    Returns:
        PIL Image object or None if generation fails
    """
    try:
        generator = AIBackgroundGenerator()
        return generator.generate_background(prompt, width, height)
    except Exception as e:
        logger.error(f"❌ Background generation failed: {e}")
        return None

def test_background_generation():
    """Test function to verify the background generator is working"""
    print("\n" + "="*60)
    print("🧪 AI BACKGROUND GENERATOR COMPATIBILITY TEST")
    print("="*60)
    
    print(f"📦 PIL Available: {'✅' if PIL_AVAILABLE else '❌'}")
    print(f"🔥 PyTorch Available: {'✅' if TORCH_AVAILABLE else '❌'}")
    print(f"🎨 Diffusers Available: {'✅' if DIFFUSERS_AVAILABLE else '❌'}")
    print(f"🔄 Force Gradient Mode: {'✅' if FORCE_GRADIENT_MODE else '❌'}")
    
    if not PIL_AVAILABLE:
        print("\n❌ CRITICAL: Cannot generate backgrounds - PIL not available")
        return False
    
    mode = "AI (Stable Diffusion)" if (DIFFUSERS_AVAILABLE and not FORCE_GRADIENT_MODE) else "Gradient Fallback"
    print(f"\n🎯 Generation Mode: {mode}")
    
    # Test multiple themes
    test_cases = [
        ("professional blue office", "Should produce blue professional gradient"),
        ("ocean sunset background", "Should produce ocean-themed gradient"),
        ("dark tech cyber background", "Should produce dark tech gradient"),
        ("green nature forest", "Should produce green nature gradient")
    ]
    
    print(f"\n🔄 Testing {len(test_cases)} different prompts...")
    
    success_count = 0
    for i, (prompt, expected) in enumerate(test_cases, 1):
        print(f"\n📋 Test {i}: '{prompt}'")
        print(f"    Expected: {expected}")
        
        try:
            image = generate_ai_background(prompt, 400, 300)
            
            if image:
                print(f"    ✅ Generated: {image.size} {image.mode} image")
                success_count += 1
                
                # Save test image
                test_path = f"test_bg_{i}.png"
                try:
                    image.save(test_path)
                    print(f"    💾 Saved: {test_path}")
                except Exception as save_error:
                    print(f"    ⚠️  Save failed: {save_error}")
            else:
                print("    ❌ Generation failed - returned None")
                
        except Exception as e:
            print(f"    ❌ Generation error: {e}")
    
    print(f"\n" + "="*60)
    print(f"📊 RESULTS: {success_count}/{len(test_cases)} tests passed")
    
    if success_count == len(test_cases):
        print("🎉 ALL TESTS PASSED! Background generator is working perfectly.")
        return True
    elif success_count > 0:
        print("⚠️  PARTIAL SUCCESS: Some backgrounds generated successfully.")
        return True
    else:
        print("❌ ALL TESTS FAILED: Background generator is not working.")
        return False

# Example usage and testing
if __name__ == "__main__":
    # Run compatibility test
    success = test_background_generation()
    
    if success:
        print(f"\n🚀 Ready to generate backgrounds!")
        print(f"💡 Usage example:")
        print(f"   from ai_background import generate_ai_background")
        print(f"   image = generate_ai_background('your prompt here')")
        print(f"   image.save('background.png')")
    else:
        print(f"\n⚠️  Please check the error messages above.")
        print(f"💡 Make sure PIL (Pillow) is installed: pip install Pillow")