import spaces
import gradio as gr
import torch
from PIL import Image
import numpy as np
from clip_interrogator import Config, Interrogator
import logging
import os
import warnings
from datetime import datetime
import gc
import re

# Suppress warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)
os.environ["TOKENIZERS_PARALLELISM"] = "false"

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

def get_device():
    if torch.cuda.is_available():
        return "cuda"
    elif torch.backends.mps.is_available():
        return "mps"
    else:
        return "cpu"

DEVICE = get_device()

class FluxRulesEngine:
    """
    Flux prompt optimization based on Pariente AI research
    Implements structured prompt generation following validated rules
    """
    
    def __init__(self):
        self.forbidden_elements = ["++", "weights", "white background [en dev]"]
        self.structure_order = {
            1: "article",
            2: "descriptive_adjectives", 
            3: "main_subject",
            4: "verb_action",
            5: "context_location",
            6: "environmental_details",
            7: "materials_textures",
            8: "lighting_effects",
            9: "technical_specs",
            10: "quality_style"
        }
        
        self.articles = ["a", "an", "the"]
        self.quality_adjectives = [
            "majestic", "pristine", "sleek", "elegant", "dramatic", 
            "cinematic", "professional", "stunning", "refined"
        ]
        
        self.lighting_types = [
            "golden hour", "studio lighting", "dramatic lighting",
            "ambient lighting", "natural light", "soft lighting",
            "rim lighting", "volumetric lighting"
        ]
        
        self.technical_specs = [
            "Shot on Phase One", "f/2.8 aperture", "50mm lens",
            "85mm lens", "35mm lens", "professional photography",
            "medium format", "high resolution"
        ]
        
        self.materials = [
            "metallic", "glass", "chrome", "leather", "fabric",
            "wood", "concrete", "steel", "ceramic"
        ]
    
    def extract_subject(self, base_prompt):
        """Extract main subject from CLIP analysis"""
        words = base_prompt.lower().split()
        
        # Common subjects to identify
        subjects = [
            "car", "vehicle", "automobile", "person", "man", "woman", 
            "building", "house", "landscape", "mountain", "tree",
            "flower", "animal", "dog", "cat", "bird"
        ]
        
        for word in words:
            if word in subjects:
                return word
                
        # Fallback to first noun-like word
        return words[0] if words else "subject"
    
    def detect_setting(self, base_prompt):
        """Detect environmental context"""
        prompt_lower = base_prompt.lower()
        
        settings = {
            "studio": ["studio", "backdrop", "seamless"],
            "outdoor": ["outdoor", "outside", "landscape", "nature"],
            "urban": ["city", "street", "urban", "building"],
            "coastal": ["beach", "ocean", "coast", "sea"],
            "indoor": ["room", "interior", "inside", "home"]
        }
        
        for setting, keywords in settings.items():
            if any(keyword in prompt_lower for keyword in keywords):
                return setting
                
        return "neutral environment"
    
    def optimize_for_flux(self, base_prompt, style_preference="professional"):
        """Apply Flux-specific optimization rules"""
        
        # Clean forbidden elements
        cleaned_prompt = base_prompt
        for forbidden in self.forbidden_elements:
            cleaned_prompt = cleaned_prompt.replace(forbidden, "")
        
        # Extract key elements
        subject = self.extract_subject(base_prompt)
        setting = self.detect_setting(base_prompt)
        
        # Build structured prompt
        components = []
        
        # 1. Article
        article = "A" if subject[0] not in 'aeiou' else "An"
        components.append(article)
        
        # 2. Descriptive adjectives (max 2-3)
        adjectives = np.random.choice(self.quality_adjectives, size=2, replace=False)
        components.extend(adjectives)
        
        # 3. Main subject
        components.append(subject)
        
        # 4. Verb/Action (gerund form)
        if "person" in subject or "man" in subject or "woman" in subject:
            action = "standing"
        else:
            action = "positioned"
        components.append(action)
        
        # 5. Context/Location
        context_map = {
            "studio": "in a professional studio setting",
            "outdoor": "in a natural outdoor environment", 
            "urban": "on an urban street",
            "coastal": "along a dramatic coastline",
            "indoor": "in an elegant interior space"
        }
        components.append(context_map.get(setting, "in a carefully composed scene"))
        
        # 6. Environmental details
        env_details = ["with subtle atmospheric effects", "surrounded by carefully balanced elements"]
        components.append(np.random.choice(env_details))
        
        # 7. Materials/Textures (if applicable)
        if any(mat in base_prompt.lower() for mat in ["car", "vehicle", "metal"]):
            material = np.random.choice(["with metallic surfaces", "featuring chrome details"])
            components.append(material)
        
        # 8. Lighting effects
        lighting = np.random.choice(self.lighting_types)
        components.append(f"illuminated by {lighting}")
        
        # 9. Technical specs
        tech_spec = np.random.choice(self.technical_specs)
        components.append(tech_spec)
        
        # 10. Quality/Style
        if style_preference == "cinematic":
            quality = "cinematic composition"
        elif style_preference == "commercial":
            quality = "commercial photography quality"
        else:
            quality = "professional photography"
            
        components.append(quality)
        
        # Join components with proper punctuation
        prompt = ", ".join(components)
        
        # Capitalize first letter
        prompt = prompt[0].upper() + prompt[1:]
        
        return prompt
    
    def get_optimization_score(self, prompt):
        """Calculate optimization score for Flux compatibility"""
        score = 0
        max_score = 100
        
        # Structure check (order compliance)
        if prompt.startswith(("A", "An", "The")):
            score += 15
            
        # Adjective count (optimal 2-3)
        adj_count = len([adj for adj in self.quality_adjectives if adj in prompt.lower()])
        if 2 <= adj_count <= 3:
            score += 15
        elif adj_count == 1:
            score += 10
            
        # Technical specs presence
        if any(spec in prompt for spec in self.technical_specs):
            score += 20
            
        # Lighting specification
        if any(light in prompt.lower() for light in self.lighting_types):
            score += 15
            
        # No forbidden elements
        if not any(forbidden in prompt for forbidden in self.forbidden_elements):
            score += 15
            
        # Proper punctuation and structure
        if "," in prompt and prompt.endswith(("photography", "composition", "quality")):
            score += 10
            
        # Length optimization (Flux works best with detailed but not excessive prompts)
        word_count = len(prompt.split())
        if 15 <= word_count <= 35:
            score += 10
        elif 10 <= word_count <= 45:
            score += 5
            
        return min(score, max_score)

class FluxPromptOptimizer:
    def __init__(self):
        self.interrogator = None
        self.flux_engine = FluxRulesEngine()
        self.usage_count = 0
        self.device = DEVICE
        self.is_initialized = False
    
    def initialize_model(self, progress_callback=None):
        if self.is_initialized:
            return True
            
        try:
            if progress_callback:
                progress_callback("Initializing CLIP model...")
            
            config = Config(
                clip_model_name="ViT-L-14/openai",
                download_cache=True,
                chunk_size=2048,
                quiet=True,
                device=self.device
            )
            
            self.interrogator = Interrogator(config)
            self.is_initialized = True
            
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
                
            return True
            
        except Exception as e:
            logger.error(f"Initialization error: {e}")
            return False
    
    def optimize_image(self, image):
        if image is None:
            return None
            
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        elif not isinstance(image, Image.Image):
            image = Image.open(image)
        
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # Optimize image size for processing
        max_size = 768 if self.device != "cpu" else 512
        if image.size[0] > max_size or image.size[1] > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        return image
    
    @spaces.GPU
    def generate_optimized_prompt(self, image, style_preference="professional", mode="best", progress_callback=None):
        try:
            if not self.is_initialized:
                if not self.initialize_model(progress_callback):
                    return "❌ Model initialization failed.", "", 0
            
            if image is None:
                return "❌ Please upload an image.", "", 0
            
            self.usage_count += 1
            
            if progress_callback:
                progress_callback("Analyzing image content...")
            
            image = self.optimize_image(image)
            if image is None:
                return "❌ Image processing failed.", "", 0
            
            if progress_callback:
                progress_callback("Extracting visual features...")
            
            start_time = datetime.now()
            
            # Get base analysis from CLIP
            try:
                if mode == "fast":
                    base_prompt = self.interrogator.interrogate_fast(image)
                elif mode == "classic":
                    base_prompt = self.interrogator.interrogate_classic(image)
                else:
                    base_prompt = self.interrogator.interrogate(image)
            except Exception as e:
                base_prompt = self.interrogator.interrogate_fast(image)
            
            if progress_callback:
                progress_callback("Applying Flux optimization rules...")
            
            # Apply Flux-specific optimization
            optimized_prompt = self.flux_engine.optimize_for_flux(base_prompt, style_preference)
            
            # Calculate optimization score
            score = self.flux_engine.get_optimization_score(optimized_prompt)
            
            end_time = datetime.now()
            duration = (end_time - start_time).total_seconds()
            
            # Memory cleanup
            if self.device == "cpu":
                gc.collect()
            else:
                torch.cuda.empty_cache()
            
            # Generate analysis info
            gpu_status = "⚡ ZeroGPU" if torch.cuda.is_available() else "💻 CPU"
            
            analysis_info = f"""
**Analysis Complete**

**Processing:** {gpu_status} • {duration:.1f}s • {mode.title()} mode  
**Style:** {style_preference.title()} photography  
**Optimization Score:** {score}/100  
**Generation:** #{self.usage_count}  

**Base Analysis:** {base_prompt[:100]}...  
**Enhancement:** Applied Flux-specific structure and terminology  
"""
            
            return optimized_prompt, analysis_info, score
            
        except Exception as e:
            return f"❌ Error: {str(e)}", "Please try with a different image or contact support.", 0

optimizer = FluxPromptOptimizer()

@spaces.GPU  
def process_image_with_progress(image, style_preference, mode):
    def progress_callback(message):
        return message
    
    yield "🚀 Initializing Flux Optimizer...", """
**Flux Prompt Optimizer**

Analyzing image with advanced computer vision  
Applying research-based optimization rules  
Generating Flux-compatible prompt structure  
""", 0
    
    prompt, info, score = optimizer.generate_optimized_prompt(image, style_preference, mode, progress_callback)
    yield prompt, info, score

def clear_outputs():
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return "", "", 0

def create_interface():
    # Professional CSS with elegant typography
    css = """
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
    
    .gradio-container {
        max-width: 1200px !important;
        margin: 0 auto !important;
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
        background: linear-gradient(135deg, #f8fafc 0%, #f1f5f9 100%) !important;
    }
    
    .main-header {
        text-align: center;
        padding: 2rem 0 3rem 0;
        background: linear-gradient(135deg, #1e293b 0%, #334155 100%);
        color: white;
        margin: -2rem -2rem 2rem -2rem;
        border-radius: 0 0 24px 24px;
    }
    
    .main-title {
        font-size: 2.5rem !important;
        font-weight: 700 !important;
        margin: 0 0 0.5rem 0 !important;
        letter-spacing: -0.025em !important;
        background: linear-gradient(135deg, #60a5fa 0%, #3b82f6 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
    }
    
    .subtitle {
        font-size: 1.125rem !important;
        font-weight: 400 !important;
        opacity: 0.8 !important;
        margin: 0 !important;
    }
    
    .section-header {
        font-size: 1.25rem !important;
        font-weight: 600 !important;
        color: #1e293b !important;
        margin: 0 0 1rem 0 !important;
        padding-bottom: 0.5rem !important;
        border-bottom: 2px solid #e2e8f0 !important;
    }
    
    .prompt-output {
        font-family: 'SF Mono', 'Monaco', 'Inconsolata', 'Roboto Mono', monospace !important;
        font-size: 14px !important;
        line-height: 1.6 !important;
        background: linear-gradient(135deg, #ffffff 0%, #f8fafc 100%) !important;
        border: 1px solid #e2e8f0 !important;
        border-radius: 12px !important;
        padding: 1.5rem !important;
        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1) !important;
    }
    
    .info-panel {
        background: linear-gradient(135deg, #f0f9ff 0%, #e0f2fe 100%) !important;
        border: 1px solid #0ea5e9 !important;
        border-radius: 12px !important;
        padding: 1.25rem !important;
        font-size: 0.875rem !important;
        line-height: 1.5 !important;
    }
    
    .score-display {
        text-align: center !important;
        padding: 1rem !important;
        background: linear-gradient(135deg, #f0fdf4 0%, #dcfce7 100%) !important;
        border: 2px solid #22c55e !important;
        border-radius: 12px !important;
        margin: 1rem 0 !important;
    }
    
    .score-number {
        font-size: 2rem !important;
        font-weight: 700 !important;
        color: #16a34a !important;
        margin: 0 !important;
    }
    
    .score-label {
        font-size: 0.875rem !important;
        color: #15803d !important;
        margin: 0 !important;
        text-transform: uppercase !important;
        letter-spacing: 0.05em !important;
    }
    """
    
    with gr.Blocks(
        theme=gr.themes.Soft(),
        title="Flux Prompt Optimizer",
        css=css
    ) as interface:
        
        gr.HTML("""
        <div class="main-header">
            <div class="main-title">⚡ Flux Prompt Optimizer</div>
            <div class="subtitle">Advanced prompt generation for Flux models • Research-based optimization</div>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("## 📷 Image Input", elem_classes=["section-header"])
                
                image_input = gr.Image(
                    label="Upload your image",
                    type="pil",
                    height=320,
                    show_label=False
                )
                
                gr.Markdown("## ⚙️ Optimization Settings", elem_classes=["section-header"])
                
                style_selector = gr.Dropdown(
                    choices=["professional", "cinematic", "commercial", "artistic"],
                    value="professional",
                    label="Photography Style",
                    info="Select the target style for prompt optimization"
                )
                
                mode_selector = gr.Dropdown(
                    choices=["fast", "classic", "best"],
                    value="best",
                    label="Analysis Mode",
                    info="Balance between speed and detail"
                )
                
                optimize_btn = gr.Button(
                    "🚀 Generate Optimized Prompt",
                    variant="primary",
                    size="lg"
                )
                
                gr.Markdown("""
                ### About Flux Optimization
                
                This tool applies research-validated rules for Flux prompt generation:
                
                • **Structured composition** following optimal element order
                • **Technical specifications** for professional results  
                • **Lighting and material** terminology optimization
                • **Quality markers** specific to Flux model architecture
                """)
            
            with gr.Column(scale=1):
                gr.Markdown("## 📝 Optimized Prompt", elem_classes=["section-header"])
                
                prompt_output = gr.Textbox(
                    label="Generated Prompt",
                    placeholder="Your optimized Flux prompt will appear here...",
                    lines=6,
                    max_lines=10,
                    elem_classes=["prompt-output"],
                    show_copy_button=True,
                    show_label=False
                )
                
                # Score display
                score_output = gr.HTML(
                    value='<div class="score-display"><div class="score-number">--</div><div class="score-label">Optimization Score</div></div>'
                )
                
                info_output = gr.Markdown(
                    value="",
                    elem_classes=["info-panel"]
                )
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear", size="sm")
                    copy_btn = gr.Button("📋 Copy Prompt", size="sm")
        
        gr.Markdown("""
        ---
        ### 🔬 Research Foundation
        
        Flux Prompt Optimizer implements validated prompt engineering research for optimal Flux model performance. 
        The optimization engine applies structured composition rules, technical terminology, and quality markers 
        specifically calibrated for Flux architecture.
        
        **Developed by Pariente AI** • Advanced AI Research Laboratory
        """)
        
        # Event handlers
        def update_score_display(score):
            color = "#22c55e" if score >= 80 else "#f59e0b" if score >= 60 else "#ef4444"
            return f'''
            <div class="score-display" style="border-color: {color};">
                <div class="score-number" style="color: {color};">{score}</div>
                <div class="score-label">Optimization Score</div>
            </div>
            '''
        
        def copy_prompt_to_clipboard(prompt):
            return prompt
        
        optimize_btn.click(
            fn=lambda img, style, mode: [
                *process_image_with_progress(img, style, mode),
                update_score_display(list(process_image_with_progress(img, style, mode))[-1][2])
            ],
            inputs=[image_input, style_selector, mode_selector],
            outputs=[prompt_output, info_output, score_output]
        )
        
        clear_btn.click(
            fn=clear_outputs,
            outputs=[prompt_output, info_output, score_output]
        )
        
        copy_btn.click(
            fn=copy_prompt_to_clipboard,
            inputs=[prompt_output],
            outputs=[]
        )
    
    return interface

if __name__ == "__main__":
    logger.info("🚀 Starting Flux Prompt Optimizer")
    interface = create_interface()
    interface.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )