# ===== CRITICAL: Import spaces FIRST =====
try:
    import spaces
    HF_SPACES = True
except ImportError:
    def spaces_gpu_decorator(duration=60):
        def decorator(func):
            return func
        return decorator
    spaces = type('spaces', (), {'GPU': spaces_gpu_decorator})()
    HF_SPACES = False

# ===== Core imports =====
import random
import os
import gradio as gr
import torch
from diffusers import DiffusionPipeline, AutoencoderTiny
from PIL import Image
import gc

# ===== Configuration =====
class Config:
    MODEL_ID = "stabilityai/sdxl-turbo"
    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
    TORCH_DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
    MAX_SEED = 2**32 - 1
    # Speed optimizations
    USE_TINY_VAE = True
    USE_XFORMERS = True if torch.cuda.is_available() else False
    USE_TORCH_COMPILE = True if torch.cuda.is_available() else False

# ===== Content Definitions =====
ADULT_STYLES = {
    "Natural": "voluptuous natural beauty, soft intimate lighting, authentic passion",
    "Artistic": "tasteful artistic nude, elegant eroticism, sensual fine art photography",
    "Sensual": "heated sensual poses, passionate romantic lighting, breathtaking chemistry",
    "Glamour": "high-end glamour photography, luxurious eroticism, stunning professional quality", 
    "Vintage": "vintage pin-up passion, retro erotic aesthetic, classic intimacy",
    "Boudoir": "intimate boudoir photography, sensual elegant setting, passionate embraces"
}

ADULT_PROMPTS = {
    "Female": [
        "breathtaking woman in passionate embrace, curves pressed against muscular male form, intimate moment captured",
        "voluptuous female arching sensually against handsome partner, hands exploring, warm glowing lighting",
        "gorgeous woman locked in passionate kiss with attractive man, bodies entwined, sensual tension",
        "beautiful lady with leg wrapped around lover's hip, intimate bedroom setting, romantic candlelight",
        "stunning female pinned against wall by muscular partner, heated moment, dramatic cinematic lighting"
    ],
    "Male": [
        "toned handsome man in sensual embrace with beautiful woman, strong hands exploring her curves",
        "muscular male pressing passionate kisses along partner's neck, intimate bedroom atmosphere",
        "attractive man lifting female partner in passionate clinch, bodies perfectly aligned",
        "fit handsome male nibbling lover's ear, sensual expression, warm golden lighting"
    ],
    "Couple": [
        "attractive couple tangled in sheets, passionate morning embrace, soft sunrise lighting",
        "gorgeous pair in heated kiss against wall, hands desperately clutching, dramatic shadows",
        "beautiful woman pressed against muscular man in shower, steam enhancing sensual moment",
        "fit couple in passionate dance-like embrace, elegant eroticism, studio lighting",
        "toned lovers in sensual kitchen encounter, barely contained passion, natural lighting"
    ]
}

# ===== Pipeline Manager =====
class UltraFastPipeline:
    def __init__(self):
        self.pipe = None
        self.loaded = False
    
    def load(self):
        if self.loaded:
            return True
        
        try:
            print("⚡ Loading optimized SDXL-Turbo...")
            
            # Load base pipeline
            self.pipe = DiffusionPipeline.from_pretrained(
                Config.MODEL_ID,
                torch_dtype=Config.TORCH_DTYPE,
                use_safetensors=True,
                variant="fp16" if Config.TORCH_DTYPE == torch.float16 else None
            )
            
            # Speed optimizations
            if Config.USE_TINY_VAE:
                self.pipe.vae = AutoencoderTiny.from_pretrained(
                    "madebyollin/taesdxl",
                    torch_dtype=Config.TORCH_DTYPE
                )
            
            if Config.USE_XFORMERS and hasattr(self.pipe, 'enable_xformers_memory_efficient_attention'):
                self.pipe.enable_xformers_memory_efficient_attention()
            
            if Config.USE_TORCH_COMPILE:
                torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True)
            
            self.pipe.to(Config.DEVICE)
            self.loaded = True
            print("✅ Model loaded with speed optimizations!")
            return True
        except Exception as e:
            print(f"❌ Failed to load model: {e}")
            return False
    
    def generate(self, prompt, negative_prompt="", width=512, height=512, steps=1, guidance=0.0, seed=None):
        if not self.loaded:
            if not self.load():
                return None
        
        if seed is None:
            seed = random.randint(0, Config.MAX_SEED)
        
        generator = torch.Generator(Config.DEVICE).manual_seed(seed)
        
        try:
            # Clear cache
            if Config.DEVICE == "cuda":
                torch.cuda.empty_cache()
            
            # Ultra-fast generation settings
            result = self.pipe(
                prompt=prompt,
                negative_prompt=negative_prompt,
                width=width,
                height=height,
                num_inference_steps=max(1, steps),  # Minimum 1 step
                guidance_scale=max(0.0, guidance),   # Can be 0 for turbo
                generator=generator,
                output_type="pil"
            ).images[0]
            
            return result, seed
        except Exception as e:
            print(f"Generation failed: {e}")
            if Config.DEVICE == "cuda":
                torch.cuda.empty_cache()
            return None, seed

# ===== Global Pipeline =====
pipeline = UltraFastPipeline()

# ===== Generation Functions =====
def get_random_prompt(prompt_type):
    return random.choice(ADULT_PROMPTS.get(prompt_type, ADULT_PROMPTS["Female"]))

@spaces.GPU(duration=20)
def generate_adult_image(
    prompt_type,
    custom_prompt,
    style,
    negative_prompt,
    width,
    height,
    steps,
    seed,
    randomize_seed
):
    try:
        # Get prompt
        base_prompt = (custom_prompt.strip() or 
                      random.choice(ADULT_PROMPTS.get(prompt_type, ADULT_PROMPTS["Female"])))
        
        # Add style
        final_prompt = f"{base_prompt}, {ADULT_STYLES.get(style, '')}" if style else base_prompt
        
        # Generate seed
        if randomize_seed:
            seed = random.randint(0, Config.MAX_SEED)
        
        # Generate image
        image, used_seed = pipeline.generate(
            prompt=final_prompt,
            negative_prompt=negative_prompt,
            width=width,
            height=height,
            steps=max(1, min(steps, 2)),  # Limit to 2 steps max for speed
            seed=seed
        )
        
        if image is None:
            return None, seed, "❌ Generation failed"
        
        return image, used_seed, f"✅ Generated in {steps} step(s)!\nSeed: {used_seed}"
        
    except Exception as e:
        return None, seed, f"❌ Error: {str(e)}"

# ===== Interface =====
def create_interface():
    with gr.Blocks(title="Ultra-Fast Content Generator") as demo:
        gr.Markdown("""
        ## ⚡ Ultra-Fast Image Content Generator
        SDXL-Turbo optimized for maximum speed (1-2 steps)
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                prompt_type = gr.Dropdown(
                    list(ADULT_PROMPTS.keys()),
                    value="Female",
                    label="Quick Prompts"
                )
                
                gr.Row([
                    gr.Button("🎲 Random Prompt", variant="secondary").click(
                        get_random_prompt, prompt_type, custom_prompt
                    ),
                    gr.Button("⚡ Generate", variant="primary").click(
                        generate_adult_image,
                        [prompt_type, custom_prompt, style, negative_prompt,
                         width, height, steps, seed, randomize_seed],
                        [result_image, seed, info_text]
                    )
                ])
                
                custom_prompt = gr.Textbox(
                    placeholder="Enter prompt...",
                    lines=3,
                    label="Custom Prompt"
                ).submit(
                    generate_adult_image,
                    [prompt_type, custom_prompt, style, negative_prompt,
                     width, height, steps, seed, randomize_seed],
                    [result_image, seed, info_text]
                )
                
                style = gr.Dropdown(
                    list(ADULT_STYLES.keys()),
                    value="Natural",
                    label="Style"
                )
                
                with gr.Row():
                    width = gr.Slider(256, 768, 512, step=64, label="Width")
                    height = gr.Slider(256, 768, 512, step=64, label="Height")
                
                with gr.Row():
                    steps = gr.Slider(1, 2, 1, step=1, label="Steps (1-2)")
                    seed = gr.Number(42, label="Seed", precision=0)
                
                randomize_seed = gr.Checkbox(True, label="Random Seed")
                
                with gr.Accordion("Advanced"):
                    negative_prompt = gr.Textbox(
                        "ugly, deformed, blurry, bad anatomy",
                        label="Negative Prompt",
                        lines=2
                    )
            
            with gr.Column(scale=1):
                result_image = gr.Image(height=500)
                info_text = gr.Textbox(lines=2, interactive=False)
        
        # Preset buttons
        gr.Row([
            gr.Button("🌸 Natural").click(
                lambda: ("Natural", get_random_prompt("Female")),
                outputs=[style, custom_prompt]
            ),
            gr.Button("🎨 Artistic").click(
                lambda: ("Artistic", get_random_prompt("Female")),
                outputs=[style, custom_prompt]
            ),
            gr.Button("✨ Glamour").click(
                lambda: ("Glamour", get_random_prompt("Female")),
                outputs=[style, custom_prompt]
            )
        ])
    
    return demo

# ===== Main =====
def main():
    print("⚡ Starting Ultra-Fast Generator...")
    demo = create_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )

if __name__ == "__main__":
    # Pre-load model when starting
    pipeline.load()
    main()