import io
import tempfile
import zipfile

import random
import torch
import spaces
import gradio as gr
from diffusers import DiffusionPipeline

MAX_SEED = 2**32 - 1


# ===== Custom aesthetic =====
# Neo-noir dusk palette with cyan + amber accents, glass panels, and subtle grain.
CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');

:root {
    /* Light Mode (Professional & Clean) */
    --bg: #fdfdfd;
    --panel: rgba(255, 255, 255, 0.95);
    --card: #ffffff;
    --border: #e5e7eb;
    --border-hover: #d1d5db;
    --text: #111827;
    --text-secondary: #4b5563;
    --muted: #9ca3af;
    --accent: #0f172a; /* Dark sleek accent for professionalism */
    --accent-hover: #1e293b;
    --accent-text: #ffffff;
    --primary-gradient: linear-gradient(135deg, #0f172a 0%, #334155 100%);
    --glow: 0 0 0 transparent;
    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.05);
    --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.05), 0 2px 4px -1px rgba(0, 0, 0, 0.03);
    --shadow-lg: 0 10px 15px -3px rgba(0, 0, 0, 0.05), 0 4px 6px -2px rgba(0, 0, 0, 0.03);
    --radius: 12px;
    --input-bg: #ffffff;
    --input-border: #e2e8f0;
    --checkbox-bg: #f1f5f9;
    --body-bg: #f8fafc; /* Very subtle cool gray */
    --font-heading: 'Inter', -apple-system, sans-serif;
    --font-body: 'Inter', -apple-system, sans-serif;
}

.dark {
    /* Dark Mode (Neo-Noir Polished) */
    --bg: #05080f;
    --panel: rgba(12, 18, 32, 0.85);
    --card: rgba(18, 28, 46, 0.70);
    --border: rgba(36, 224, 194, 0.15);
    --border-hover: rgba(36, 224, 194, 0.3);
    --text: #e9f3ff;
    --text-secondary: #94a3b8;
    --muted: #64748b;
    --accent: #24e0c2;
    --accent-hover: #18cdb0;
    --accent-text: #041019;
    --primary-gradient: linear-gradient(120deg, #24e0c2 0%, #ffb347 100%);
    --glow: 0 8px 32px rgba(36, 224, 194, 0.12);
    --shadow-sm: 0 1px 2px 0 rgba(0, 0, 0, 0.2);
    --shadow-md: 0 4px 6px -1px rgba(0, 0, 0, 0.3);
    --shadow-lg: 0 20px 40px -5px rgba(0, 0, 0, 0.4);
    --radius: 16px;
    --input-bg: rgba(255,255,255,0.03);
    --input-border: rgba(255,255,255,0.08);
    --checkbox-bg: #0d1829;
    --body-bg: radial-gradient(circle at 20% 20%, rgba(36, 224, 194, 0.06), transparent 35%),
               radial-gradient(circle at 82% 12%, rgba(0, 156, 196, 0.06), transparent 35%),
               linear-gradient(145deg, #05080f 0%, #080f1e 100%);
    --font-heading: 'Inter', -apple-system, sans-serif;
    --font-body: 'Inter', -apple-system, sans-serif;
}

body, .gradio-container {
    font-family: var(--font-body) !important;
    background: var(--body-bg) !important;
    color: var(--text);
    min-height: 100vh;
}

/* Titles & Typography */
.gradio-container .prose h1,
.gradio-container .prose h2,
.gradio-container .prose h3 {
    font-family: var(--font-heading);
    letter-spacing: -0.025em;
    font-weight: 700;
    color: var(--text);
}

.gradio-container .prose h1 {
    font-size: 2.25rem;
    margin-bottom: 0.5rem;
    background: var(--primary-gradient);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    display: inline-block;
}

.gradio-container * { letter-spacing: -0.01em; }

/* Panels & Cards */
.gr-block, .gr-panel, .gr-group {
    background: var(--panel);
    border: 1px solid var(--border);
    border-radius: var(--radius);
    box-shadow: var(--shadow-sm);
    backdrop-filter: blur(8px);
    transition: box-shadow 0.2s ease, border-color 0.2s ease;
}

.hero-card {
    background: var(--card);
    border: 1px solid var(--border);
    padding: 24px;
    border-radius: var(--radius);
    box-shadow: var(--shadow-md);
    position: relative;
    overflow: hidden;
}

.tagline {
    display: inline-flex;
    align-items: center;
    gap: 8px;
    padding: 6px 14px;
    background: var(--input-bg);
    border: 1px solid var(--border);
    border-radius: 999px;
    font-size: 0.875rem;
    font-weight: 500;
    color: var(--text-secondary);
    margin-bottom: 12px;
}

.hero-card p {
    color: var(--text-secondary);
    font-size: 1.05rem;
    line-height: 1.6;
    max-width: 65ch;
}

/* Inputs */
textarea, input:not([type='checkbox']):not([type='radio']),
.gr-input, .gr-textbox, .gr-number, .gr-slider input {
    background: var(--input-bg) !important;
    border: 1px solid var(--input-border) !important;
    border-radius: 10px !important;
    color: var(--text) !important;
    font-family: var(--font-body);
    transition: all 0.2s ease;
}

textarea:focus, input:focus, .gr-input:focus-within {
    border-color: var(--text-secondary) !important;
    box-shadow: 0 0 0 2px rgba(var(--accent), 0.1);
}

label, .gr-box label {
    color: var(--text-secondary) !important;
    font-weight: 600;
    font-size: 0.875rem;
    margin-bottom: 6px;
    text-transform: none !important;
}

/* Sliders */
.gr-slider input[type='range'] {
    accent-color: var(--accent);
}

/* Buttons */
.gr-button-primary, button.primary {
    background: var(--primary-gradient) !important;
    color: var(--accent-text) !important;
    font-weight: 600 !important;
    border: 1px solid rgba(255,255,255,0.1) !important;
    box-shadow: var(--shadow-md);
    border-radius: 10px !important;
    padding: 10px 24px;
    transition: transform 0.1s, box-shadow 0.2s;
}

.gr-button-primary:hover {
    transform: translateY(-1px);
    box-shadow: var(--shadow-lg);
    filter: brightness(1.1);
}

.gr-button-secondary, button.secondary, .gr-downloadbutton {
    background: var(--input-bg) !important;
    border: 1px solid var(--border) !important;
    color: var(--text) !important;
    font-weight: 500;
    border-radius: 10px !important;
    box-shadow: var(--shadow-sm);
}
.gr-button-secondary:hover {
    border-color: var(--border-hover) !important;
    background: var(--card) !important;
}
.gr-downloadbutton, .gr-downloadbutton > button { width: 100%; }

/* Gallery */
.gr-gallery {
    background: var(--input-bg);
    border-radius: var(--radius);
    border: 1px solid var(--border);
    padding: 8px;
}
.gr-gallery .thumbnail-item {
    border-radius: 8px;
    overflow: hidden;
    box-shadow: var(--shadow-sm);
    border: 1px solid transparent;
    transition: all 0.2s;
}
.gr-gallery .thumbnail-item:hover {
    box-shadow: var(--shadow-md);
    transform: scale(1.02);
}
.gr-gallery img { object-fit: cover; }

/* Footer */
.footer-note {
    color: var(--muted);
    font-size: 0.875rem;
    text-align: center;
    margin-top: 2rem;
    opacity: 0.8;
}
.footer-note a {
    color: var(--text-secondary);
    text-decoration: none;
    border-bottom: 1px dotted var(--muted);
}
.footer-note a:hover {
    color: var(--accent);
    border-bottom-style: solid;
}
"""

# Load the pipeline once at startup
print("Loading Z-Image-Turbo pipeline...")
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)
pipe.to("cuda")

'# ======== AoTI compilation + FA3 ======== (disabled on HF to avoid outdated AOTI/FA3 package errors)'
# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")

print("Pipeline loaded!")

@spaces.GPU
def generate_image(
    prompt,
    negative_prompt,
    height,
    width,
    images_count,
    num_inference_steps,
    guidance_scale,
    seed,
    randomize_seed,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate N images using a deterministic seed cascade (x1..xN)."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    base_seed = int(seed) % MAX_SEED
    if base_seed < 0:
        base_seed += MAX_SEED

    # Cap to prevent excessive VRAM usage / latency spikes on the demo space
    images_count = max(1, min(int(images_count), 12))

    seeds = [(base_seed * i) % MAX_SEED for i in range(1, images_count + 1)]

    neg_prompt = None
    if isinstance(negative_prompt, str) and negative_prompt.strip():
        neg_prompt = negative_prompt

    images = []
    image_paths = []
    for s in seeds:
        generator = torch.Generator("cuda").manual_seed(int(s))
        image = pipe(
            prompt=prompt,
            negative_prompt=neg_prompt,
            height=int(height),
            width=int(width),
            num_inference_steps=int(num_inference_steps),
            guidance_scale=float(guidance_scale),  # 0.0 is recommended default for Turbo
            generator=generator,
        ).images[0]
        images.append(image)
        tmp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        image.save(tmp_img.name, format="PNG")
        image_paths.append(tmp_img.name)

    return images, ", ".join(str(s) for s in seeds), image_paths, base_seed



def append_history(new_images, history):
    """Append new images to the history state."""
    if history is None:
        history = []
    updated_history = history + new_images
    return updated_history, updated_history


def package_zip(image_paths):
    """Pack the current image list into a ZIP file for download."""
    if not image_paths:
        raise gr.Error("No images in history to download.")

    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
    with zipfile.ZipFile(tmp, "w", zipfile.ZIP_DEFLATED) as zf:
        for idx, path in enumerate(image_paths, start=1):
            # Store as image_001.png, image_002.png, ...
            zf.write(path, arcname=f"image_{idx:03d}.png")

    tmp.flush()
    return tmp.name


# Example prompts
examples = [
    ["Astronaut riding a horse on Mars, cinematic lighting, sci-fi concept art, highly detailed"],
    ["Portrait of a wise old wizard with a long white beard, holding a glowing crystal staff, magical forest background"],
]

# Build the Gradio interface
# Build the Gradio interface
with gr.Blocks(title="Z-Image-Turbo Demo", css=CUSTOM_CSS, analytics_enabled=False) as demo:
    image_state = gr.State([])
    history_state = gr.State([])
    gr.Markdown(
        """
        <div class="hero-card">
          <div class="tagline">⚡ Turbo diffusion · 8 steps · CUDA ready</div>
          <h1>Z‑Image Turbo Studio</h1>
          <p>Draft up to twelve stylized candidates in one pass. Neo‑noir gradients, glass panels, and crisp typography keep the tooling out of your way while you explore ideas.</p>
        </div>
        """,
        sanitize_html=False,
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="e.g. bioluminescent reef city at dusk, cinematic, anamorphic glow",
                lines=4,
            )

            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="noise, blur, extra limbs, text watermark",
                lines=3,
            )

            with gr.Row():
                height = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Height",
                )
                width = gr.Slider(
                    minimum=512,
                    maximum=2048,
                    value=1024,
                    step=64,
                    label="Width",
                )

            with gr.Row():
                num_inference_steps = gr.Slider(
                    minimum=1,
                    maximum=20,
                    value=9,
                    step=1,
                    label="Inference Steps",
                    info="9 steps → 8 DiT forwards",
                )

            images_count = gr.Slider(
                minimum=1,
                maximum=12,
                value=4,
                step=1,
                label="Images",
                info="1–12 (higher counts use more VRAM)",
            )

            guidance_scale = gr.Slider(
                minimum=0.0,
                maximum=7.0,
                value=0.0,
                step=0.1,
                label="CFG Guidance Scale",
                info="0 = no CFG (recommended for Turbo models)",
            )
            
            with gr.Row():
                seed = gr.Number(
                    label="Base Seed",
                    value=42,
                    precision=0,
                )
                randomize_seed = gr.Checkbox(
                    label="Randomize",
                    value=True,
                    interactive=True,
                )
            
            generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")
        
        with gr.Column(scale=1):
            output_images = gr.Gallery(
                label="Generated Grid",
                columns=4,
                rows=None,
                preview=True,
            )
            used_seeds = gr.Textbox(
                label="Seed Cascade (x1 · x2 · ... · xN)",
                interactive=False,
            )
            history_gallery = gr.Gallery(
                label="History",
                columns=6,
                rows=None,
                preview=True,
                object_fit="cover"
            )
            download_btn = gr.DownloadButton(
                label="📦 Download All History (ZIP)",
            )
    
    gr.Markdown("### 💡 Quick Prompts")
    gr.Examples(
        examples=examples,
        inputs=[prompt],
        cache_examples=False,
    )

    gr.Markdown(
        """
        <div class="footer-note">
        Model: Tongyi-MAI/Z-Image-Turbo (Apache 2.0). Demo by <a href="https://z-image-turbo.tech" target="_blank">https://z-image-turbo.tech</a>
        </div>
        """,
        sanitize_html=False,
    )
    
    # Connect the generate button
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, height, width, images_count, num_inference_steps, guidance_scale, seed, randomize_seed],
        outputs=[output_images, used_seeds, image_state, seed],
    ).success(
        fn=append_history,
        inputs=[image_state, history_state],
        outputs=[history_state, history_gallery],
    )
    
    # Also allow generating by pressing Enter in the prompt box
    prompt.submit(
        fn=generate_image,
        inputs=[prompt, negative_prompt, height, width, images_count, num_inference_steps, guidance_scale, seed, randomize_seed],
        outputs=[output_images, used_seeds, image_state, seed],
    ).success(
        fn=append_history,
        inputs=[image_state, history_state],
        outputs=[history_state, history_gallery],
    )

    download_btn.click(
        fn=package_zip,
        inputs=[history_state],
        outputs=[download_btn],
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)