# ===== 必须首先导入spaces =====
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 os
from datetime import datetime
import random
import torch
import gradio as gr
from diffusers import AutoPipelineForText2Image, FlowMatchEulerDiscreteScheduler
from PIL import Image
import traceback
import numpy as np
import gc

# ===== 配置常量 =====
COMPEL_AVAILABLE = False
print("⚠️ Compel disabled for FLUX compatibility")

STYLE_PRESETS = {
    "None": "",
    "Realistic": "photorealistic, 8k, ultra-detailed, cinematic lighting, masterpiece",
    "Anime": "anime style, detailed, high quality, masterpiece, best quality",
    "Comic": "comic book style, bold outlines, vibrant colors, cel shading",
    "Watercolor": "watercolor illustration, soft gradients, pastel palette"
}

FIXED_MODEL = "aoxo/flux.1dev-abliterated"

QUALITY_ENHANCERS = [
    "detailed anatomy", "perfect anatomy", "soft skin",
    "high resolution", "masterpiece", "best quality",
    "professional photography", "natural lighting"
]

STYLE_ENHANCERS = {
    "Realistic": ["photorealistic", "ultra realistic", "natural lighting"],
    "Anime": ["anime style", "high quality anime", "detailed eyes"],
    "Comic": ["comic book style", "bold outlines", "vibrant colors"],
    "Watercolor": ["watercolor style", "artistic", "soft gradients"]
}

SAVE_DIR = "generated_images"
os.makedirs(SAVE_DIR, exist_ok=True)

# ===== 全局变量 =====
pipeline = None
device = None
model_loaded = False


# ===== 工具函数(必须在装饰器之前定义) =====
def cleanup_memory():
    """清理GPU内存"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.synchronize()
    gc.collect()


def enhance_prompt(prompt: str, style: str) -> str:
    """增强提示词"""
    quality_terms = ", ".join(QUALITY_ENHANCERS[:3])
    
    style_terms = ""
    if style in STYLE_ENHANCERS:
        style_terms = ", " + ", ".join(STYLE_ENHANCERS[style][:2])
    
    style_suffix = STYLE_PRESETS.get(style, "")
    
    enhanced_parts = [prompt.strip()]
    
    if style_suffix:
        enhanced_parts.append(style_suffix)
    
    if style_terms:
        enhanced_parts.append(style_terms.lstrip(", "))
    
    enhanced_parts.append(quality_terms)
    
    enhanced_prompt = ", ".join(filter(None, enhanced_parts))
    
    if len(enhanced_prompt) > 800:
        enhanced_prompt = enhanced_prompt[:800]
    
    return enhanced_prompt


def create_metadata_content(prompt, enhanced_prompt, seed, steps, cfg_scale, width, height, style):
    """创建元数据内容"""
    timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    return f"""Generated Image Metadata
======================
Timestamp: {timestamp}
Original Prompt: {prompt}
Enhanced Prompt: {enhanced_prompt}
Seed: {seed}
Steps: {steps}
CFG Scale: {cfg_scale}
Dimensions: {width}x{height}
Style: {style}
Model: FLUX.1-dev
"""


# ===== 装饰器定义(必须在使用之前) =====
def apply_spaces_decorator(func):
    """应用 spaces 装饰器,增加更长的超时时间"""
    if SPACES_AVAILABLE:
        return spaces.GPU(duration=120)(func)
    return func


# ===== 模型相关函数 =====
def initialize_model():
    """优化的模型初始化函数"""
    global pipeline, device, model_loaded
    
    if model_loaded and pipeline is not None:
        return True
    
    try:
        cleanup_memory()
        
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"🖥️ Using device: {device}")
        
        print(f"📦 Loading fixed model: {FIXED_MODEL}")
        
        pipeline = AutoPipelineForText2Image.from_pretrained(
            FIXED_MODEL,
            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
            variant=None,
            use_safetensors=True
        )
        
        pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
            pipeline.scheduler.config
        )
        pipeline = pipeline.to(device)
        
        if torch.cuda.is_available():
            # 关键优化:使用sequential代替model cpu offload
            pipeline.enable_sequential_cpu_offload()
            pipeline.enable_vae_slicing()
            pipeline.enable_vae_tiling()
        
        print("✅ Model initialization complete (Optimized)")
        model_loaded = True
        return True
        
    except Exception as e:
        print(f"❌ Critical model loading error: {e}")
        print(traceback.format_exc())
        cleanup_memory()
        model_loaded = False
        return False


@apply_spaces_decorator
def generate_image(prompt: str, style: str, negative_prompt: str = "", 
                  steps: int = 15, cfg_scale: float = 3.5, 
                  seed: int = -1, width: int = 1024, height: int = 1024, 
                  progress=gr.Progress()):
    """图像生成函数(优化版本)"""
    try:
        if not prompt or prompt.strip() == "":
            return None, "", "❌ Please enter a prompt"
        
        # 优化的参数限制
        steps = max(10, min(steps, 25))
        width = min(width, 1024)
        height = min(height, 1024)
        
        progress(0.1, desc="Initializing model...")
        if not initialize_model():
            cleanup_memory()
            return None, "", "❌ Failed to initialize model"
        
        progress(0.2, desc="Processing prompt...")
        
        if seed == -1:
            seed = random.randint(0, np.iinfo(np.int32).max)
        
        enhanced_prompt = enhance_prompt(prompt.strip(), style)
        
        if not negative_prompt.strip():
            negative_prompt = "(low quality, worst quality:1.4), (bad anatomy, bad hands:1.2), blurry, deformed"
        
        generator = torch.Generator("cpu").manual_seed(seed)
        
        progress(0.4, desc="Starting generation...")
        print(f"🔥 Inference: steps={steps}, guidance={cfg_scale}, size={width}x{height}")
        
        cleanup_memory()
        
        # 关键改动:提高max_sequence_length
        with torch.no_grad():
            result = pipeline(
                prompt=enhanced_prompt,
                negative_prompt=negative_prompt,
                num_inference_steps=steps,
                guidance_scale=cfg_scale,
                width=width,
                height=height,
                max_sequence_length=512,  # 从256改到512
                generator=generator,
                output_type="pil"
            )
        
        image = result.images[0]
        print("✅ Inference complete")
        
        progress(0.9, desc="Finalizing...")
        
        del result
        cleanup_memory()
        
        filename = f"IMG_{seed}.png"
        filepath = os.path.join(SAVE_DIR, filename)
        image.save(filepath, format="PNG", optimize=True)
        
        metadata_content = create_metadata_content(
            prompt, enhanced_prompt, seed, steps, cfg_scale, width, height, style
        )
        
        progress(1.0, desc="Complete!")
        
        generation_info = f"Prompt: {prompt}\nSeed: {seed} | Size: {width}×{height} | Steps: {steps} | CFG: {cfg_scale}"
        
        return image, generation_info, metadata_content
        
    except torch.cuda.OutOfMemoryError as e:
        cleanup_memory()
        error_msg = "❌ GPU memory insufficient. Try 768x768 or fewer steps."
        print(f"CUDA OOM: {error_msg}")
        return None, "", error_msg
        
    except Exception as e:
        cleanup_memory()
        error_msg = str(e)
        print(f"❌ Generation error: {error_msg}")
        print(traceback.format_exc())
        return None, "", f"❌ Generation failed: {error_msg}"


# ===== CSS样式 =====
css = """
/* 保持原有CSS不变 */
.gradio-container {
    max-width: 100% !important;
    margin: 0 !important;
    padding: 0 !important;
    background: linear-gradient(135deg, #e6a4f2 0%, #1197e4 100%) !important;
    min-height: 100vh !important;
}

.main-content {
    background: rgba(255, 255, 255, 0.95) !important;
    border-radius: 20px !important;
    padding: 20px !important;
    margin: 15px !important;
    box-shadow: 0 10px 25px rgba(0,0,0,0.2) !important;
}

.title {
    text-align: center !important;
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    -webkit-background-clip: text !important;
    -webkit-text-fill-color: transparent !important;
    font-size: 2rem !important;
    margin-bottom: 15px !important;
    font-weight: bold !important;
}

.generate-btn {
    background: linear-gradient(45deg, #bb6ded, #08676b) !important;
    color: white !important;
    border: none !important;
    padding: 15px 25px !important;
    border-radius: 25px !important;
    font-size: 16px !important;
    font-weight: bold !important;
    width: 100% !important;
}
"""


# ===== 创建UI =====
def create_interface():
    with gr.Blocks(css=css, title="NSFW FLUX Image Generator") as interface:
        with gr.Column(elem_classes=["main-content"]):
            gr.HTML('<div class="title">NSFW FLUX Image Generator</div>')
            
            gr.HTML('<div class="warning-box">⚠️ 18+ CONTENT WARNING ⚠️</div>')
            
            with gr.Row():
                with gr.Column(scale=2):
                    prompt_input = gr.Textbox(
                        label="Main Prompt",
                        placeholder="beautiful woman, detailed portrait...",
                        lines=6,
                        elem_classes=["prompt-box"]
                    )
                    
                    gr.HTML('''
                        <div style="background: rgba(255, 193, 7, 0.1); padding: 10px; border-radius: 8px; margin: 10px 0;">
                            <small><strong>💡 Optimized Settings:</strong><br>
                            • Max sequence length: 512 tokens (supports longer prompts)<br>
                            • Recommended steps: 15-20 for best speed/quality<br>
                            • Try 768x768 for faster generation</small>
                        </div>
                    ''')
                    
                    negative_prompt_input = gr.Textbox(
                        label="Negative Prompt (Optional)",
                        placeholder="low quality, blurry...",
                        lines=3,
                        elem_classes=["prompt-box"]
                    )
                
                with gr.Column(scale=1):
                    with gr.Group():
                        style_input = gr.Radio(
                            label="Style Preset",
                            choices=list(STYLE_PRESETS.keys()),
                            value="Realistic"
                        )
                    
                    with gr.Group():
                        seed_input = gr.Number(
                            label="Seed (-1 for random)",
                            value=-1,
                            precision=0
                        )
                    
                    with gr.Group():
                        size_preset = gr.Radio(
                            label="Size (smaller = faster)",
                            choices=["768x768", "1024x1024"],
                            value="768x768"
                        )
                    
                    with gr.Group():
                        steps_input = gr.Slider(
                            label="Steps (15-20 recommended)",
                            minimum=10,
                            maximum=25,
                            value=15,
                            step=1
                        )
                        
                        cfg_input = gr.Slider(
                            label="CFG Scale",
                            minimum=1.0,
                            maximum=15.0,
                            value=3.5,
                            step=0.1
                        )
                    
                    generate_button = gr.Button(
                        "GENERATE",
                        elem_classes=["generate-btn"],
                        variant="primary"
                    )
            
            image_output = gr.Image(
                label="Generated Image",
                elem_classes=["image-output"],
                show_label=False
            )
            
            generation_info = gr.Textbox(
                label="Generation Info",
                interactive=False,
                visible=False
            )
            
            metadata_content = gr.Textbox(visible=False)
            current_seed = gr.Number(visible=False)
            current_image = gr.Image(visible=False)
            
            with gr.Row(visible=False) as download_row:
                download_image_btn = gr.Button("Save Image", size="sm")
                download_metadata_btn = gr.Button("Save Metadata", size="sm")
        
        def parse_size(size_str):
            """解析尺寸字符串"""
            size = int(size_str.split('x')[0])
            return size, size
        
        def on_generate(prompt, style, neg_prompt, steps, cfg, seed, size_preset):
            width, height = parse_size(size_preset)
            image, info, metadata = generate_image(
                prompt, style, neg_prompt, steps, cfg, seed, width, height
            )
            
            if image is not None:
                try:
                    actual_seed = seed if seed != -1 else int(info.split("Seed:")[1].split("|")[0].strip())
                except:
                    actual_seed = seed if seed != -1 else random.randint(0, 999999)
                
                return (
                    image, info, metadata, actual_seed, image,
                    gr.update(visible=True), gr.update(visible=True)
                )
            else:
                return (
                    None, info, "", 0, None,
                    gr.update(visible=False), gr.update(visible=False)
                )
        
        generate_button.click(
            fn=on_generate,
            inputs=[
                prompt_input, style_input, negative_prompt_input, 
                steps_input, cfg_input, seed_input, size_preset
            ],
            outputs=[
                image_output, generation_info, metadata_content, 
                current_seed, current_image, generation_info, download_row
            ],
            show_progress=True
        )
        
        prompt_input.submit(
            fn=on_generate,
            inputs=[
                prompt_input, style_input, negative_prompt_input, 
                steps_input, cfg_input, seed_input, size_preset
            ],
            outputs=[
                image_output, generation_info, metadata_content, 
                current_seed, current_image, generation_info, download_row
            ],
            show_progress=True
        )
    
    return interface


# ===== 启动应用 =====
if __name__ == "__main__":
    print("🎨 Starting NSFW FLUX Image Generator (Optimized)...")
    print(f"🔧 Fixed Model: {FIXED_MODEL}")
    print(f"🔧 CUDA: {'✅ Available' if torch.cuda.is_available() else '❌ Not Available'}")
    
    app = create_interface()
    app.queue(max_size=5, default_concurrency_limit=1)
    
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=False
    )