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

# Global variable to store the pipeline
pipe = None

def load_model():
    """
    Load the Z-Image Turbo model before inference.
    This ensures the model is downloaded and ready before any generation requests.
    """
    global pipe
    
    if pipe is not None:
        return pipe
    
    print("Loading Z-Image Turbo model...")
    print("This may take a few minutes on first run while the model downloads...")
    
    try:
        # Load the pipeline with optimal settings
        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)
        print(f"Model loaded on {device}")
        
        # Optional: Enable Flash Attention for better efficiency
        try:
            pipe.transformer.set_attention_backend("flash")
            print("Flash Attention enabled")
        except Exception as e:
            print(f"Flash Attention not available: {e}")
            print("Using default attention backend")
        
        print("Model loaded successfully!")
        return pipe
        
    except Exception as e:
        print(f"Error loading model: {e}")
        raise

# Pre-load the model when the app starts
print("Initializing model on startup...")
try:
    load_model()
    print("Model initialization complete!")
except Exception as e:
    print(f"Warning: Could not pre-load model: {e}")
    print("Model will be loaded on first generation request")

@spaces.GPU()
def generate_image(
    prompt,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Generate an image using Z-Image Turbo model.
    
    Args:
        prompt: Text description of the desired image
    
    Returns:
        Generated PIL Image
    """
    global pipe
    
    # Ensure model is loaded
    if pipe is None:
        progress(0, desc="Loading model...")
        load_model()
    
    if not prompt.strip():
        raise gr.Error("Please enter a prompt to generate an image.")
    
    # Determine device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    
    # Set random seed for reproducibility
    generator = torch.Generator(device).manual_seed(42)
    
    # Generate the image with optimal settings
    progress(0.1, desc="Generating image...")
    
    try:
        result = pipe(
            prompt=prompt,
            negative_prompt=None,
            height=1024,
            width=1024,
            num_inference_steps=9,
            guidance_scale=0.0,
            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 minimal 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: 900px !important;
    margin: 0 auto !important;
    background: #ffffff !important;
    padding: 0 !important;
}

body {
    background: #f5f5f7 !important;
}

.main-container {
    background: white;
    min-height: 100vh;
    padding: 3rem 2rem;
}

.header-section {
    text-align: center;
    margin-bottom: 3rem;
    padding-bottom: 2rem;
    border-bottom: 1px solid #e5e5e7;
}

.header-section h1 {
    font-size: 3rem !important;
    font-weight: 600 !important;
    color: #1d1d1f !important;
    margin: 0 0 0.5rem 0 !important;
    letter-spacing: -1px;
}

.header-section p {
    font-size: 1.25rem !important;
    color: #86868b !important;
    margin: 0 !important;
    font-weight: 400;
}

.attribution {
    margin-top: 1rem;
    font-size: 0.875rem;
    color: #86868b;
}

.attribution a {
    color: #0071e3 !important;
    text-decoration: none;
    font-weight: 500;
    transition: opacity 0.2s;
}

.attribution a:hover {
    opacity: 0.7;
}

.prompt-section {
    margin-bottom: 2rem;
}

textarea {
    border: 1px solid #d2d2d7 !important;
    border-radius: 12px !important;
    padding: 1rem !important;
    font-size: 1.0625rem !important;
    line-height: 1.5 !important;
    transition: all 0.2s ease !important;
    background: #fbfbfd !important;
}

textarea:focus {
    border-color: #0071e3 !important;
    box-shadow: 0 0 0 4px rgba(0, 113, 227, 0.1) !important;
    outline: none !important;
    background: white !important;
}

textarea::placeholder {
    color: #86868b !important;
}

.generate-btn {
    background: #0071e3 !important;
    border: none !important;
    border-radius: 12px !important;
    padding: 0.875rem 2.5rem !important;
    font-size: 1.0625rem !important;
    font-weight: 500 !important;
    color: white !important;
    transition: all 0.2s ease !important;
    box-shadow: 0 2px 8px rgba(0, 113, 227, 0.2) !important;
    width: 100% !important;
    margin-bottom: 2rem !important;
}

.generate-btn:hover {
    background: #0077ed !important;
    box-shadow: 0 4px 12px rgba(0, 113, 227, 0.3) !important;
    transform: translateY(-1px) !important;
}

.generate-btn:active {
    transform: translateY(0) !important;
}

.output-section {
    margin-top: 2rem;
    border-radius: 12px;
    overflow: hidden;
    background: #fbfbfd;
    border: 1px solid #e5e5e7;
}

.output-section img {
    border-radius: 12px;
    width: 100%;
    height: auto;
    display: block;
}

label {
    display: none !important;
}

.footer-note {
    text-align: center;
    padding: 2rem 1rem;
    color: #86868b;
    font-size: 0.875rem;
    margin-top: 3rem;
    border-top: 1px solid #e5e5e7;
}

.footer-note p {
    margin: 0.25rem 0;
}

/* Hide unnecessary Gradio elements */
.gradio-container .footer {
    display: none !important;
}

@media (max-width: 768px) {
    .main-container {
        padding: 2rem 1rem;
    }
    
    .header-section h1 {
        font-size: 2rem !important;
    }
    
    .header-section p {
        font-size: 1rem !important;
    }
}
"""

# Create the Gradio interface
with gr.Blocks(
    title="Z-Image Turbo",
    fill_height=False,
    css=apple_css
) as demo:
    
    with gr.Column(elem_classes="main-container"):
        # Header
        gr.HTML("""
            <div class="header-section">
                <h1>Z-Image Turbo</h1>
                <p>Create stunning images from text</p>
                <div class="attribution">
                    Built with <a href="https://huggingface.co/spaces/akhaliq/anycoder" target="_blank">anycoder</a>
                </div>
            </div>
        """)
        
        # Prompt input
        with gr.Column(elem_classes="prompt-section"):
            prompt = gr.Textbox(
                placeholder="Describe the image you want to create...",
                lines=3,
                max_lines=6,
                show_label=False
            )
        
        # Generate button
        generate_btn = gr.Button(
            "Generate",
            elem_classes="generate-btn",
            size="lg"
        )
        
        # Output image
        with gr.Column(elem_classes="output-section"):
            output_image = gr.Image(
                type="pil",
                show_label=False,
                show_download_button=True,
                show_share_button=False,
                container=False
            )
        
        # Footer
        gr.HTML("""
            <div class="footer-note">
                <p>Powered by Z-Image Turbo from Tongyi-MAI</p>
                <p>Optimized for fast, high-quality image generation</p>
            </div>
        """)
    
    # Event handlers
    generate_btn.click(
        fn=generate_image,
        inputs=prompt,
        outputs=output_image,
        api_name="generate"
    )
    
    # Also allow generation on Enter key
    prompt.submit(
        fn=generate_image,
        inputs=prompt,
        outputs=output_image
    )

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