try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ Spaces available - ZeroGPU mode")
except ImportError:
    SPACES_AVAILABLE = False
    print("⚠️ Spaces not available - running in regular mode")

import gradio as gr
import torch
from diffusers import DiffusionPipeline, StableDiffusionXLPipeline
from PIL import Image
import datetime
import io
import json
import os
import re
from typing import Optional, List, Dict
import numpy as np

# ======================
# Configuration Section (Modify here to expand)
# ======================

# 1. Base Model - Using reliable SDXL models (avoiding incomplete "bait" models)
BASE_MODELS = {
    "nsfw":"votepurchase/novaFurryXL_illustriousV7b",
    "sdxl_base": "stabilityai/stable-diffusion-xl-base-1.0",  # Most reliable choice
    "realistic_vision": "SG161222/RealVisXL_V4.0",  # High-quality realistic model
    "anime_xl": "Linaqruf/animagine-xl-3.1",  # Popular anime-style SDXL
    "juggernaut_xl": "RunDiffusion/Juggernaut-XL-v9",  # High-quality general purpose
    "playground_v2": "playgroundai/playground-v2.5-1024px-aesthetic"  # Aesthetic focused
}

# Current model selection (change this to switch models)
CURRENT_MODEL_KEY = "nsfw"  # Changed to working model
BASE_MODEL = BASE_MODELS[CURRENT_MODEL_KEY]

# 2. Fixed LoRAs (Auto-loaded, not user-selectable) - Using actual LoRA models
FIXED_LORAS = {
    "detail_enhancer": {
        "repo_id": "ostris/ikea-instructions-lora-sdxl",  # Real LoRA for details
        "filename": None,
        "weight": 0.6,
        "trigger_words": "high quality, detailed, sharp focus"
    },
    "quality_boost": {
        "repo_id": "stabilityai/stable-diffusion-xl-offset-example-lora",  # Official SDXL LoRA
        "filename": None,
        "weight": 0.5,
        "trigger_words": "masterpiece, best quality"
    }
}

# 3. Style Templates (Auto-prepended to user prompts)
STYLE_PROMPTS = {
    "None": "",
    "Realistic": "photorealistic, ultra-detailed skin, natural lighting, 8k uhd, professional photography, DSLR, soft lighting, high quality, film grain, Fujifilm XT3, masterpiece, ",
    "Anime": "anime style, cel shading, vibrant colors, detailed eyes, studio ghibli style, manga style, trending on pixiv, masterpiece, ",
    "Comic": "comic book style, bold outlines, dynamic angles, comic panel, Marvel DC style, inked lines, pop art, masterpiece, ",
    "Watercolor": "watercolor painting, soft brush strokes, translucent layers, artistic, painterly, paper texture, traditional art, masterpiece, ",
}

# 4. Optional LoRAs (User-selectable via dropdown, can select multiple) - Using real, verified LoRAs
OPTIONAL_LORAS = {
    "None": {
        "repo_id": None,
        "weight": 0.0,
        "trigger_words": "",
        "description": "No additional LoRA"
    },
    "Offset Noise LoRA": {
        "repo_id": "stabilityai/stable-diffusion-xl-offset-example-lora",
        "weight": 0.7,
        "trigger_words": "high contrast, dramatic lighting",
        "description": "Enhanced contrast and lighting (Official Stability AI)"
    },
    "LCM LoRA": {
        "repo_id": "latent-consistency/lcm-lora-sdxl",
        "weight": 0.8,
        "trigger_words": "lcm style, high quality",
        "description": "Latent Consistency Model for faster generation"
    },
    "Pixel Art LoRA": {
        "repo_id": "nerijs/pixel-art-xl",
        "weight": 0.9,
        "trigger_words": "pixel art style, 8bit, retro game",
        "description": "Pixel art style transformation"
    },
    "Watercolor LoRA": {
        "repo_id": "ostris/watercolor-style-lora-sdxl",
        "weight": 0.8,
        "trigger_words": "watercolor painting, soft colors, artistic",
        "description": "Watercolor painting style"
    },
    "Sketch LoRA": {
        "repo_id": "ostris/crayon-style-lora-sdxl",
        "weight": 0.7,
        "trigger_words": "sketch style, pencil drawing, artistic",
        "description": "Hand-drawn sketch style"
    },
    "Portrait LoRA": {
        "repo_id": "ostris/face-helper-sdxl-lora",
        "weight": 0.8,
        "trigger_words": "portrait, beautiful face, detailed eyes",
        "description": "Portrait and face enhancement"
    }
}

# Default Parameters
DEFAULT_SEED = -1
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_LORA_SCALE = 0.8
DEFAULT_STEPS = 30
DEFAULT_CFG = 7.5

# Supported Languages (for future expansion)
SUPPORTED_LANGUAGES = {
    "en": "English",
    "zh": "中文",
    "ja": "日本語",
    "ko": "한국어"
}

# ======================
# Global Variables: Lazy Loading
# ======================
pipe = None
current_loras = {}
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_pipeline():
    """Load the base Illustrious XL pipeline with fallback options"""
    global pipe
    if pipe is None:
        print(f"🚀 Loading base model: {BASE_MODEL}...")
        
        # Try to load the selected model with fallback options
        model_loaded = False
        models_to_try = [BASE_MODEL]
        
        # Add fallback models if primary fails
        if CURRENT_MODEL_KEY != "sdxl_base":
            models_to_try.append(BASE_MODELS["sdxl_base"])
        if CURRENT_MODEL_KEY != "realistic_vision":
            models_to_try.append(BASE_MODELS["realistic_vision"])
            
        for model_id in models_to_try:
            try:
                print(f"Attempting to load: {model_id}")
                pipe = StableDiffusionXLPipeline.from_pretrained(
                    model_id,
                    torch_dtype=torch.float16,
                    use_safetensors=True,
                    variant="fp16"
                ).to(device)
                
                # Enable memory optimizations for ZeroGPU
                pipe.enable_attention_slicing()
                pipe.enable_vae_slicing()
                if hasattr(pipe, 'enable_model_cpu_offload'):
                    pipe.enable_model_cpu_offload()
                if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
                    pipe.enable_xformers_memory_efficient_attention()
                
                print(f"✅ Successfully loaded: {model_id}")
                model_loaded = True
                break
                
            except Exception as e:
                print(f"❌ Failed to load {model_id}: {e}")
                continue
        
        if not model_loaded:
            raise Exception("Failed to load any model. Please check your configuration.")
            
    return pipe

def unload_pipeline():
    """Unload pipeline to free memory"""
    global pipe, current_loras
    if pipe is not None:
        # Clear any loaded LoRAs
        try:
            pipe.unload_lora_weights()
        except:
            pass
        del pipe
        torch.cuda.empty_cache()
        pipe = None
        current_loras = {}
        print("🗑️ Pipeline unloaded.")

def load_lora_weights(lora_configs: List[Dict]):
    """Load multiple LoRA weights efficiently with error handling"""
    global pipe, current_loras
    
    if not lora_configs:
        return
    
    # Unload existing LoRAs if different
    new_lora_ids = [config['repo_id'] for config in lora_configs if config['repo_id']]
    if set(current_loras.keys()) != set(new_lora_ids):
        try:
            pipe.unload_lora_weights()
            current_loras = {}
        except:
            pass
    
    # Load new LoRAs with better error handling
    adapter_names = []
    adapter_weights = []
    
    for config in lora_configs:
        if config['repo_id'] and config['repo_id'] not in current_loras:
            try:
                # Try different loading methods
                adapter_name = config['name'].replace(' ', '_').lower()
                
                # Method 1: Direct loading
                pipe.load_lora_weights(
                    config['repo_id'], 
                    adapter_name=adapter_name
                )
                current_loras[config['repo_id']] = adapter_name
                print(f"✅ Loaded LoRA: {config['name']}")
                
            except Exception as e:
                print(f"⚠️ Failed to load LoRA {config['name']}: {e}")
                # Skip this LoRA and continue with others
                continue
        
        # Add to active adapters if successfully loaded
        if config['repo_id'] in current_loras:
            adapter_names.append(current_loras[config['repo_id']])
            adapter_weights.append(config['weight'])
    
    # Set adapter weights if any adapters loaded
    if adapter_names:
        try:
            pipe.set_adapters(adapter_names, adapter_weights=adapter_weights)
            print(f"✅ Activated {len(adapter_names)} LoRA adapters")
        except Exception as e:
            print(f"⚠️ Warning setting adapter weights: {e}")
            # Try without weights
            try:
                pipe.set_adapters(adapter_names)
            except:
                print("❌ Failed to set any adapters")

def process_long_prompt(prompt: str, max_length: int = 77) -> str:
    """Process long prompts by intelligent truncation and optimization"""
    if len(prompt.split()) <= max_length:
        return prompt
    
    # Split into sentences and prioritize
    sentences = re.split(r'[.!?]+', prompt)
    sentences = [s.strip() for s in sentences if s.strip()]
    
    # Keep most important parts (first sentence + key descriptors)
    if sentences:
        result = sentences[0]
        remaining = max_length - len(result.split())
        
        for sentence in sentences[1:]:
            words = sentence.split()
            if len(words) <= remaining:
                result += ". " + sentence
                remaining -= len(words)
            else:
                # Add partial sentence with most important words
                important_words = [w for w in words if len(w) > 3][:remaining]
                if important_words:
                    result += ". " + " ".join(important_words)
                break
        
        return result
    
    return " ".join(prompt.split()[:max_length])

# ======================
# Main Generation Function
# ======================
@spaces.GPU(duration=60) if SPACES_AVAILABLE else lambda x: x
def generate_image(
    prompt: str,
    negative_prompt: str,
    style: str,
    seed: int,
    width: int,
    height: int,
    selected_loras: List[str],
    lora_scale: float,
    steps: int,
    cfg_scale: float,
    language: str = "en"
):
    """Main image generation function with ZeroGPU optimization"""
    global pipe
    
    try:
        # Load pipeline
        pipe = load_pipeline()
        
        # Handle seed
        if seed == -1:
            seed = torch.randint(0, 2**32, (1,)).item()
        generator = torch.Generator(device=device).manual_seed(seed)
        
        # Process prompts
        style_prefix = STYLE_PROMPTS.get(style, "")
        processed_prompt = process_long_prompt(style_prefix + prompt, max_length=150)
        processed_negative = process_long_prompt(negative_prompt, max_length=100)
        
        # Prepare LoRA configurations
        lora_configs = []
        active_trigger_words = []
        
        # Add fixed LoRAs
        for name, config in FIXED_LORAS.items():
            if config["repo_id"]:
                lora_configs.append({
                    'name': name,
                    'repo_id': config["repo_id"],
                    'weight': config["weight"]
                })
                if config["trigger_words"]:
                    active_trigger_words.append(config["trigger_words"])
        
        # Add selected optional LoRAs
        for lora_name in selected_loras:
            if lora_name != "None" and lora_name in OPTIONAL_LORAS:
                config = OPTIONAL_LORAS[lora_name]
                if config["repo_id"]:
                    lora_configs.append({
                        'name': lora_name,
                        'repo_id': config["repo_id"],
                        'weight': config["weight"] * lora_scale
                    })
                    if config["trigger_words"]:
                        active_trigger_words.append(config["trigger_words"])
        
        # Load LoRAs
        load_lora_weights(lora_configs)
        
        # Combine trigger words with prompt
        if active_trigger_words:
            trigger_text = ", ".join(active_trigger_words)
            final_prompt = f"{processed_prompt}, {trigger_text}"
        else:
            final_prompt = processed_prompt
        
        # Generate image
        with torch.autocast(device):
            image = pipe(
                prompt=final_prompt,
                negative_prompt=processed_negative,
                num_inference_steps=steps,
                guidance_scale=cfg_scale,
                width=width,
                height=height,
                generator=generator,
            ).images[0]
        
        # Generate metadata
        timestamp = datetime.datetime.now()
        metadata = {
            "prompt": final_prompt,
            "original_prompt": prompt,
            "negative_prompt": processed_negative,
            "base_model": BASE_MODEL,
            "style": style,
            "fixed_loras": [name for name in FIXED_LORAS.keys()],
            "selected_loras": [name for name in selected_loras if name != "None"],
            "lora_scale": lora_scale,
            "seed": seed,
            "steps": steps,
            "cfg_scale": cfg_scale,
            "width": width,
            "height": height,
            "language": language,
            "timestamp": timestamp.isoformat(),
            "trigger_words": active_trigger_words
        }
        
        # Generate filenames
        timestamp_str = timestamp.strftime("%y%m%d%H%M")
        filename_base = f"{seed}-{timestamp_str}"
        
        # Save image as WebP
        img_buffer = io.BytesIO()
        image.save(img_buffer, format="WEBP", quality=95, method=6)
        img_buffer.seek(0)
        
        # Save metadata as JSON
        metadata_str = json.dumps(metadata, indent=2, ensure_ascii=False)
        
        return (
            image,
            metadata_str
        )
        
    except Exception as e:
        error_msg = f"Generation failed: {str(e)}"
        print(f"❌ {error_msg}")
        return None, error_msg

# ======================
# Gradio Interface
# ======================
def create_interface():
    """Create the Gradio interface"""
    
    with gr.Blocks(
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="green",
            neutral_hue="slate",
        ).set(
            body_background_fill="linear-gradient(135deg, #1e40af, #059669)",
            button_primary_background_fill="white",
            button_primary_text_color="#1e40af",
            input_background_fill="rgba(255,255,255,0.9)",
            block_background_fill="rgba(255,255,255,0.1)",
        ),
        css="""
        body { 
            font-family: 'Segoe UI', 'Arial', sans-serif; 
            background: linear-gradient(135deg, #1e40af, #059669);
        }
        .gr-button { 
            font-family: 'Segoe UI', 'Arial', sans-serif; 
            font-weight: 600;
            border-radius: 8px;
        }
        .gr-textbox { 
            font-family: 'Consolas', 'Monaco', 'Courier New', monospace; 
            border-radius: 8px;
        }
        .gr-dropdown, .gr-slider, .gr-radio { 
            border-radius: 8px;
        }
        .gr-form { 
            background: rgba(255,255,255,0.05);
            border-radius: 16px;
            padding: 20px;
            margin: 10px;
        }
        """,
        title="AI Photo Generator - Illustrious XL"
    ) as demo:
        
        gr.Markdown("""
        # 🎨 AI Photo Generator (Illustrious XL + Multi-LoRA)
        """)
        
        with gr.Row():
            # Left Column - Controls
            with gr.Column(scale=3, elem_classes=["gr-form"]):
                
                # a. Prompt Input
                prompt_input = gr.Textbox(
                    label="Prompt (Positive)",
                    placeholder="A beautiful woman with flowing hair, golden hour lighting, cinematic composition, high detail...",
                    lines=6,
                    max_lines=20,
                    elem_classes=["gr-textbox"]
                )
                
                # b. Negative Prompt Input
                negative_prompt_input = gr.Textbox(
                    label="Negative Prompt",
                    value="blurry, low quality, deformed, cartoon, anime, text, watermark, signature, username, worst quality, low res, bad anatomy, bad hands, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, bad feet, extra fingers, mutated hands, poorly drawn hands, bad proportions, extra limbs, disfigured, ugly, gross proportions, malformed limbs",
                    lines=4,
                    max_lines=15,
                    elem_classes=["gr-textbox"]
                )
                
                # c. Style Selection
                style_radio = gr.Radio(
                    choices=list(STYLE_PROMPTS.keys()),
                    label="Style Template",
                    value="Realistic",
                    elem_classes=["gr-radio"]
                )
                
                # Multi-row controls
                with gr.Row():
                    # d. Seed Control
                    with gr.Column():
                        seed_input = gr.Slider(
                            minimum=-1,
                            maximum=99999999,
                            step=1,
                            value=DEFAULT_SEED,
                            label="Seed (-1 = Random)"
                        )
                        seed_reset = gr.Button("Reset Seed", size="sm")
                
                with gr.Row():
                    # e. Width Control
                    with gr.Column():
                        width_input = gr.Slider(
                            minimum=512,
                            maximum=1536,
                            step=64,
                            value=DEFAULT_WIDTH,
                            label="Width"
                        )
                        width_reset = gr.Button("Reset Width", size="sm")
                    
                    # f. Height Control
                    with gr.Column():
                        height_input = gr.Slider(
                            minimum=512,
                            maximum=1536,
                            step=64,
                            value=DEFAULT_HEIGHT,
                            label="Height"
                        )
                        height_reset = gr.Button("Reset Height", size="sm")
                
                # g. LoRA Selection (Multi-select)
                lora_dropdown = gr.Dropdown(
                    choices=list(OPTIONAL_LORAS.keys()),
                    label="Optional LoRAs (Multi-select)",
                    value=["None"],
                    multiselect=True,
                    elem_classes=["gr-dropdown"]
                )
                
                # h. LoRA Scale Control
                with gr.Row():
                    lora_scale_slider = gr.Slider(
                        minimum=0.0,
                        maximum=1.5,
                        step=0.05,
                        value=DEFAULT_LORA_SCALE,
                        label="LoRA Scale"
                    )
                    lora_reset = gr.Button("Reset LoRA", size="sm")
                
                # i. Generation Controls
                with gr.Row():
                    steps_slider = gr.Slider(
                        minimum=10,
                        maximum=100,
                        step=1,
                        value=DEFAULT_STEPS,
                        label="Steps"
                    )
                    cfg_slider = gr.Slider(
                        minimum=1.0,
                        maximum=20.0,
                        step=0.5,
                        value=DEFAULT_CFG,
                        label="CFG Scale"
                    )
                    gen_reset = gr.Button("Reset Generation", size="sm")
                
                # Language Selection (Optional)
                language_dropdown = gr.Dropdown(
                    choices=list(SUPPORTED_LANGUAGES.keys()),
                    label="Language (Optional)",
                    value="en",
                    visible=True  # Hidden for now, can be enabled later
                )
                
                # m. Generate Button
                generate_btn = gr.Button(
                    "✨ Generate Image", 
                    variant="primary", 
                    size="lg",
                    elem_classes=["gr-button"]
                )
            
            # Right Column - Outputs
            with gr.Column(scale=2):
                # j. Image Display
                image_output = gr.Image(
                    label="Generated Image", 
                    height=600, 
                    format="webp"
                )
                
                # Simplified UI without complex download buttons
                with gr.Row():
                    gr.Markdown("**Right-click the image above to download**")
                
                # k. Metadata Display
                metadata_output = gr.Textbox(
                    label="Generation Metadata (JSON)",
                    lines=15,
                    max_lines=25,
                    elem_classes=["gr-textbox"]
                )
        
        # ======================
        # Event Handlers
        # ======================
        
        # Reset buttons
        seed_reset.click(fn=lambda: -1, outputs=seed_input)
        width_reset.click(fn=lambda: DEFAULT_WIDTH, outputs=width_input)
        height_reset.click(fn=lambda: DEFAULT_HEIGHT, outputs=height_input)
        lora_reset.click(fn=lambda: DEFAULT_LORA_SCALE, outputs=lora_scale_slider)
        gen_reset.click(
            fn=lambda: (DEFAULT_STEPS, DEFAULT_CFG),
            outputs=[steps_slider, cfg_slider]
        )
        
        # Main generation function
        def generate_and_prepare_downloads(*args):
            result = generate_image(*args)
            if result[0] is not None:  # Success
                image, metadata, img_filename, meta_filename = result
                
                # Save files temporarily for download
                import tempfile
                import os
                
                # Create temporary files
                temp_dir = tempfile.mkdtemp()
                img_path = os.path.join(temp_dir, img_filename)
                meta_path = os.path.join(temp_dir, meta_filename)
                
                # Save image
                image.save(img_path, format="WEBP", quality=95)
                
                # Save metadata
                with open(meta_path, 'w', encoding='utf-8') as f:
                    f.write(metadata)
                
                return (
                    image,
                    metadata,
                    img_path,   # File path for download
                    meta_path   # File path for download
                )
            else:  # Error
                return result[0], result[1], None, None
        
        # Generate button click - Simplified without complex downloads
        generate_btn.click(
            fn=generate_image,
            inputs=[
                prompt_input, negative_prompt_input, style_radio,
                seed_input, width_input, height_input,
                lora_dropdown, lora_scale_slider,
                steps_slider, cfg_slider, language_dropdown
            ],
            outputs=[
                image_output, metadata_output
            ]
        )
        
        # Show LoRA descriptions
        def show_lora_info(selected_loras):
            if not selected_loras or selected_loras == ["None"]:
                return "No LoRAs selected"
            
            info = "Selected LoRAs:\n"
            for lora_name in selected_loras:
                if lora_name in OPTIONAL_LORAS:
                    config = OPTIONAL_LORAS[lora_name]
                    info += f"• {lora_name}: {config['description']}\n"
                    if config['trigger_words']:
                        info += f"  Triggers: {config['trigger_words']}\n"
            return info
        
        lora_dropdown.change(
            fn=show_lora_info,
            inputs=[lora_dropdown],
            outputs=[gr.Textbox(label="LoRA Information", visible=False)]
        )
    
    return demo

# ======================
# Launch Application
# ======================
if __name__ == "__main__":
    demo = create_interface()
    demo.queue(max_size=20)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )