import gradio as gr
import torch
from diffusers import ZImagePipeline
import os
from pathlib import Path
import spaces

# Initialize the pipeline
print("Loading Z-Image Turbo model...")
pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)

# Move to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe.to(device)

# Optional: Enable Flash Attention for better efficiency
try:
    pipe.transformer.set_attention_backend("flash")
    print("Flash Attention enabled")
except:
    print("Flash Attention not available, using default")

print("Model loaded successfully!")

@spaces.GPU()
def generate_image(
    prompt,
    negative_prompt,
    height,
    width,
    num_steps,
    seed,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generate an image using Z-Image Turbo model.
    
    Args:
        prompt: Text description of the desired image
        negative_prompt: What to avoid in the image
        height: Image height
        width: Image width
        num_steps: Number of inference steps
        seed: Random seed for reproducibility
    
    Returns:
        Generated PIL Image
    """
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to generate an image.")
    
    # Set random seed for reproducibility
    generator = torch.Generator(device).manual_seed(int(seed))
    
    # Generate the image
    progress(0, desc="Initializing generation...")
    
    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt if negative_prompt.strip() else None,
            height=int(height),
            width=int(width),
            num_inference_steps=int(num_steps),
            guidance_scale=0.0,  # Turbo models use 0 guidance
            generator=generator,
        )
        
        image = result.images[0]
        progress(1.0, desc="Complete!")
        
        return image
    
    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")

# Apple-inspired CSS
apple_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');

* {
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif !important;
}

.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    background: linear-gradient(135deg, #f5f7fa 0%, #e8ecf1 100%) !important;
}

.main-header {
    text-align: center;
    padding: 2.5rem 1rem 1.5rem 1rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 24px;
    margin-bottom: 2rem;
    box-shadow: 0 10px 40px rgba(102, 126, 234, 0.3);
}

.main-header h1 {
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    color: white !important;
    margin: 0 0 0.5rem 0 !important;
    letter-spacing: -0.5px;
}

.main-header p {
    font-size: 1.1rem !important;
    color: rgba(255, 255, 255, 0.9) !important;
    margin: 0 !important;
    font-weight: 400;
}

.attribution {
    text-align: center;
    margin-top: 0.5rem;
    font-size: 0.9rem;
    color: rgba(255, 255, 255, 0.8);
}

.attribution a {
    color: white !important;
    text-decoration: none;
    font-weight: 600;
    transition: opacity 0.2s;
}

.attribution a:hover {
    opacity: 0.8;
}

.control-panel {
    background: white;
    border-radius: 20px;
    padding: 1.5rem;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
    margin-bottom: 1.5rem;
}

.output-panel {
    background: white;
    border-radius: 20px;
    padding: 1.5rem;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.08);
    min-height: 600px;
}

.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 0.875rem 2rem !important;
    font-size: 1rem !important;
    font-weight: 600 !important;
    color: white !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.3) !important;
}

.primary-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.4) !important;
}

.secondary-btn {
    background: #f3f4f6 !important;
    border: 1px solid #e5e7eb !important;
    border-radius: 12px !important;
    padding: 0.875rem 2rem !important;
    font-size: 1rem !important;
    font-weight: 500 !important;
    color: #374151 !important;
    transition: all 0.3s ease !important;
}

.secondary-btn:hover {
    background: #e5e7eb !important;
    transform: translateY(-1px) !important;
}

textarea, input, .input-container {
    border-radius: 12px !important;
    border: 1.5px solid #e5e7eb !important;
    transition: all 0.2s ease !important;
    font-size: 0.95rem !important;
}

textarea:focus, input:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}

.slider-container input[type="range"] {
    accent-color: #667eea !important;
}

label {
    font-weight: 600 !important;
    color: #1f2937 !important;
    font-size: 0.9rem !important;
    margin-bottom: 0.5rem !important;
}

.example-card {
    border-radius: 12px !important;
    transition: all 0.3s ease !important;
    cursor: pointer !important;
    border: 2px solid transparent !important;
}

.example-card:hover {
    border-color: #667eea !important;
    transform: scale(1.02) !important;
}

.accordion {
    border-radius: 12px !important;
    border: 1.5px solid #e5e7eb !important;
    overflow: hidden !important;
}

#output-image {
    border-radius: 16px !important;
    overflow: hidden !important;
}

.footer-note {
    text-align: center;
    padding: 1.5rem;
    color: #6b7280;
    font-size: 0.9rem;
    margin-top: 2rem;
}

@media (max-width: 768px) {
    .main-header h1 {
        font-size: 1.75rem !important;
    }
    
    .main-header p {
        font-size: 0.95rem !important;
    }
    
    .control-panel, .output-panel {
        padding: 1rem;
    }
}
"""

# Example prompts
examples = [
    [
        "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights.",
        "",
        1024,
        1024,
        9,
        42
    ],
    [
        "A serene Japanese garden with cherry blossoms in full bloom, koi pond with crystal clear water, traditional wooden bridge, soft morning light filtering through trees, ultra detailed, photorealistic",
        "blurry, low quality, distorted",
        1024,
        1024,
        9,
        123
    ],
    [
        "Futuristic cyberpunk city at night, neon signs reflecting on wet streets, flying cars, towering skyscrapers with holographic advertisements, rain, cinematic lighting, highly detailed",
        "daytime, bright, low quality",
        1024,
        1024,
        9,
        456
    ],
    [
        "Majestic dragon soaring through clouds at sunset, scales shimmering with iridescent colors, mountains in background, fantasy art style, epic composition, dramatic lighting",
        "cartoon, simple, flat",
        1024,
        1024,
        9,
        789
    ],
    [
        "Cozy coffee shop interior, warm lighting, wooden furniture, plants on shelves, barista preparing coffee, steam rising from cup, bokeh background, inviting atmosphere",
        "empty, cold, harsh lighting",
        1024,
        1024,
        9,
        321
    ]
]

# Create the Gradio interface
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="violet",
        secondary_hue="purple",
        neutral_hue="slate",
        font=gr.themes.GoogleFont("Inter")
    ),
    css=apple_css,
    title="Z-Image Turbo - AI Image Generator",
    fill_height=False
) as demo:
    
    # Header
    with gr.Row():
        with gr.Column():
            gr.HTML("""
                <div class="main-header">
                    <h1>✨ Z-Image Turbo</h1>
                    <p>Create stunning images from text in seconds</p>
                    <div class="attribution">
                        Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">anycoder</a>
                    </div>
                </div>
            """)
    
    with gr.Row():
        # Left column - Controls
        with gr.Column(scale=1, elem_classes="control-panel"):
            prompt = gr.Textbox(
                label="✍️ Prompt",
                placeholder="Describe the image you want to create...",
                lines=4,
                max_lines=8
            )
            
            negative_prompt = gr.Textbox(
                label="🚫 Negative Prompt (Optional)",
                placeholder="What you don't want in the image...",
                lines=2,
                max_lines=4
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False, elem_classes="accordion"):
                with gr.Row():
                    width = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="Width",
                        elem_classes="slider-container"
                    )
                    height = gr.Slider(
                        minimum=512,
                        maximum=2048,
                        value=1024,
                        step=64,
                        label="Height",
                        elem_classes="slider-container"
                    )
                
                num_steps = gr.Slider(
                    minimum=4,
                    maximum=20,
                    value=9,
                    step=1,
                    label="Inference Steps",
                    info="More steps = better quality but slower",
                    elem_classes="slider-container"
                )
                
                seed = gr.Number(
                    label="Seed",
                    value=42,
                    precision=0,
                    info="Use same seed for reproducible results"
                )
                
                random_seed_btn = gr.Button(
                    "🎲 Random Seed",
                    elem_classes="secondary-btn",
                    size="sm"
                )
            
            with gr.Row():
                generate_btn = gr.Button(
                    "🎨 Generate Image",
                    variant="primary",
                    elem_classes="primary-btn",
                    size="lg"
                )
                clear_btn = gr.ClearButton(
                    [prompt, negative_prompt],
                    value="🗑️ Clear",
                    elem_classes="secondary-btn",
                    size="lg"
                )
        
        # Right column - Output
        with gr.Column(scale=1, elem_classes="output-panel"):
            output_image = gr.Image(
                label="Generated Image",
                type="pil",
                elem_id="output-image",
                show_label=False,
                show_download_button=True,
                show_share_button=True,
                height=600
            )
    
    # Examples section
    with gr.Row():
        gr.Examples(
            examples=examples,
            inputs=[prompt, negative_prompt, height, width, num_steps, seed],
            outputs=output_image,
            fn=generate_image,
            cache_examples=False,
            label="💡 Try these examples",
            examples_per_page=5,
            elem_classes="example-card"
        )
    
    # Footer
    gr.HTML("""
        <div class="footer-note">
            <p>Powered by <strong>Z-Image Turbo</strong> from Tongyi-MAI | 
            Optimized for fast, high-quality image generation</p>
            <p style="margin-top: 0.5rem; font-size: 0.85rem;">
                💡 Tip: Be specific and detailed in your prompts for best results
            </p>
        </div>
    """)
    
    # Event handlers
    def randomize_seed():
        import random
        return random.randint(0, 2**32 - 1)
    
    random_seed_btn.click(
        fn=randomize_seed,
        inputs=[],
        outputs=seed
    )
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, height, width, num_steps, seed],
        outputs=output_image,
        api_name="generate"
    )
    
    # Also allow generation on Enter key in prompt
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, negative_prompt, height, width, num_steps, seed],
        outputs=output_image
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        share=False,
        show_error=True,
        favicon_path=None
    )