"""
The-X — AI Image Generation & Editing
Powered by FLUX.1-schnell (Text-to-Image) & FLUX.1-Kontext-dev (Image Editing)
"""

import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
from PIL import Image
from diffusers import FluxPipeline, FluxKontextPipeline

# ─── Load Models ───────────────────────────────────────────────────────────────

# Text-to-Image pipeline
t2i_pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=torch.bfloat16,
).to("cuda")

# Image Editing pipeline
edit_pipe = FluxKontextPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-Kontext-dev",
    torch_dtype=torch.bfloat16,
).to("cuda")

MAX_SEED = np.iinfo(np.int32).max

# ─── Text-to-Image Function ────────────────────────────────────────────────────

@spaces.GPU(duration=120)
def generate_image(
    prompt: str,
    seed: int,
    randomize_seed: bool,
    width: int,
    height: int,
    num_inference_steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate an image from a text prompt using FLUX.1-schnell."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator().manual_seed(seed)
    image = t2i_pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0,
    ).images[0]
    return image, seed


# ─── Image Editing Function ────────────────────────────────────────────────────

@spaces.GPU(duration=120)
def edit_image(
    input_image,
    prompt: str,
    seed: int,
    randomize_seed: bool,
    guidance_scale: float,
    num_inference_steps: int,
    progress=gr.Progress(track_tqdm=True),
):
    """Edit an image using FLUX.1-Kontext-dev based on a text instruction."""
    if input_image is None:
        raise gr.Error("Please upload an image to edit.")
    if not prompt.strip():
        raise gr.Error("Please enter an editing instruction.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    input_image = input_image.convert("RGB")
    generator = torch.Generator().manual_seed(seed)

    image = edit_pipe(
        image=input_image,
        prompt=prompt,
        guidance_scale=guidance_scale,
        width=input_image.size[0],
        height=input_image.size[1],
        num_inference_steps=num_inference_steps,
        generator=generator,
    ).images[0]
    return image, seed


# ─── UI ─────────────────────────────────────────────────────────────────────────

CSS = """
/* The-X Branding */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700;800;900&display=swap');

:root {
    --bg-primary: #0a0a0f;
    --bg-secondary: #12121a;
    --bg-card: #1a1a2e;
    --accent: #6c63ff;
    --accent-glow: #8b83ff;
    --text-primary: #ffffff;
    --text-secondary: #a0a0b8;
    --border: #2a2a3e;
}

body {
    background: var(--bg-primary) !important;
    color: var(--text-primary) !important;
    font-family: 'Inter', sans-serif !important;
}

#the-x-header {
    text-align: center;
    padding: 2rem 1rem 1rem;
    background: linear-gradient(180deg, rgba(108,99,255,0.15) 0%, transparent 100%);
    border-bottom: 1px solid var(--border);
    margin-bottom: 1rem;
}

#the-x-header h1 {
    font-size: 3.5rem !important;
    font-weight: 900 !important;
    background: linear-gradient(135deg, #6c63ff, #ff6b9d, #ffd93d);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    margin: 0 !important;
    letter-spacing: -2px;
}

#the-x-header p {
    color: var(--text-secondary) !important;
    font-size: 1.1rem !important;
    margin-top: 0.5rem !important;
}

.gradio-container {
    max-width: 1200px !important;
    background: var(--bg-primary) !important;
}

.tabs {
    background: var(--bg-secondary) !important;
    border-radius: 12px !important;
    border: 1px solid var(--border) !important;
}

.tab-nav button {
    color: var(--text-secondary) !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    padding: 12px 24px !important;
    border-radius: 8px !important;
}

.tab-nav button.selected {
    background: var(--accent) !important;
    color: white !important;
}

.gr-input, .gr-text-input textarea {
    background: var(--bg-card) !important;
    border: 1px solid var(--border) !important;
    color: var(--text-primary) !important;
    border-radius: 10px !important;
    font-size: 1rem !important;
}

.gr-button-primary {
    background: linear-gradient(135deg, var(--accent), var(--accent-glow)) !important;
    border: none !important;
    border-radius: 10px !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
    padding: 12px 32px !important;
    color: white !important;
    box-shadow: 0 4px 15px rgba(108,99,255,0.3) !important;
    transition: all 0.3s ease !important;
}

.gr-button-primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 25px rgba(108,99,255,0.5) !important;
}

.gr-accordion {
    background: var(--bg-card) !important;
    border: 1px solid var(--border) !important;
    border-radius: 10px !important;
}

.gr-slider input {
    accent-color: var(--accent) !important;
}

.gr-image {
    border-radius: 12px !important;
    border: 1px solid var(--border) !important;
}

.gr-examples {
    background: var(--bg-secondary) !important;
    border-radius: 12px !important;
    border: 1px solid var(--border) !important;
    padding: 1rem !important;
}

.gr-examples button {
    background: var(--bg-card) !important;
    border: 1px solid var(--border) !important;
    color: var(--text-primary) !important;
    border-radius: 8px !important;
}

#col-container {
    margin: 0 auto;
    max-width: 1100px;
}

.model-badge {
    display: inline-block;
    padding: 4px 12px;
    border-radius: 20px;
    font-size: 0.75rem;
    font-weight: 600;
    margin: 0 4px;
}

.badge-gen {
    background: rgba(108,99,255,0.2);
    color: #8b83ff;
    border: 1px solid rgba(108,99,255,0.3);
}

.badge-edit {
    background: rgba(255,107,157,0.2);
    color: #ff6b9d;
    border: 1px solid rgba(255,107,157,0.3);
}

.badge-quality {
    background: rgba(255,217,61,0.2);
    color: #ffd93d;
    border: 1px solid rgba(255,217,61,0.3);
}

.footer-info {
    text-align: center;
    padding: 2rem 0;
    color: var(--text-secondary);
    border-top: 1px solid var(--border);
    margin-top: 2rem;
}
"""

with gr.Blocks(css=CSS, title="The-X — AI Image Generation & Editing") as demo:

    # ─── Header ────────────────────────────────────────────────────────────
    with gr.Column(elem_id="the-x-header"):
        gr.Markdown(
            """
            # The-X
            **Next-Generation AI Image Generation & Editing**
            
            <span class="model-badge badge-gen">FLUX.1-schnell</span>
            <span class="model-badge badge-edit">FLUX.1-Kontext</span>
            <span class="model-badge badge-quality">8K Hyper-Detailed</span>
            """
        )

    with gr.Tabs() as tabs:

        # ══════════════════════════════════════════════════════════════════════
        # TAB 1: TEXT-TO-IMAGE
        # ══════════════════════════════════════════════════════════════════════
        with gr.Tab("🎨 Generate Image", id="t2i"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown(
                    "### Create images from text descriptions\n"
                    "Powered by **FLUX.1-schnell** — a 12B parameter rectified flow transformer "
                    "for ultra-fast, high-quality generation."
                )

                with gr.Row():
                    with gr.Column(scale=3):
                        prompt_t2i = gr.Textbox(
                            label="Prompt",
                            placeholder="Describe the image you want to create in detail...",
                            lines=3,
                            max_lines=8,
                        )
                        with gr.Row():
                            run_t2i = gr.Button("✨ Generate", variant="primary", scale=1)

                    with gr.Column(scale=2):
                        result_t2i = gr.Image(
                            label="Generated Image",
                            show_label=True,
                            interactive=False,
                            height=512,
                        )

                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    with gr.Row():
                        seed_t2i = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=0,
                        )
                        randomize_seed_t2i = gr.Checkbox(
                            label="Randomize seed", value=True
                        )

                    with gr.Row():
                        width_t2i = gr.Slider(
                            label="Width",
                            minimum=512,
                            maximum=2048,
                            step=32,
                            value=1024,
                        )
                        height_t2i = gr.Slider(
                            label="Height",
                            minimum=512,
                            maximum=2048,
                            step=32,
                            value=1024,
                        )

                    num_steps_t2i = gr.Slider(
                        label="Inference Steps",
                        minimum=1,
                        maximum=12,
                        step=1,
                        value=4,
                        info="More steps = better quality (4 is optimal for schnell)",
                    )

                gr.Examples(
                    examples=[
                        [
                            "A hyper-detailed portrait of a cyberpunk samurai standing on a neon-lit rooftop in Tokyo at night, rain pouring, holographic advertisements reflecting in puddles, 8K resolution, cinematic lighting, photorealistic"
                        ],
                        [
                            "An ancient dragon made of crystalline ice perched atop a mountain peak during a blizzard, aurora borealis in the sky, intricate ice crystal formations, ultra-detailed fantasy art, 8K"
                        ],
                        [
                            "A futuristic space station orbiting Saturn, with detailed ring system visible through panoramic windows, astronauts floating in zero gravity, photorealistic, volumetric lighting, 8K render"
                        ],
                        [
                            "A microscopic view of a bioluminescent deep-sea creature, glowing tentacles in the abyss, particles of marine snow, extreme macro photography, 8K detail, National Geographic style"
                        ],
                        [
                            "A steampunk clockwork owl perched on a brass branch, intricate gears and cogs visible through translucent enamel, warm candlelight glow, hyper-detailed, 8K"
                        ],
                    ],
                    inputs=[prompt_t2i],
                    label="✨ Example Prompts — Click to try",
                )

                # Send to edit tab button
                send_to_edit = gr.Button("📤 Send to Image Editor", visible=False)

        # ══════════════════════════════════════════════════════════════════════
        # TAB 2: IMAGE EDITING
        # ══════════════════════════════════════════════════════════════════════
        with gr.Tab("✏️ Edit Image", id="edit"):
            with gr.Column(elem_id="col-container"):
                gr.Markdown(
                    "### Transform images with text instructions\n"
                    "Powered by **FLUX.1-Kontext-dev** — state-of-the-art instruction-based "
                    "image editing with character consistency."
                )

                with gr.Row():
                    with gr.Column(scale=2):
                        input_image_edit = gr.Image(
                            label="📷 Upload Image to Edit",
                            type="pil",
                            height=400,
                        )

                    with gr.Column(scale=2):
                        result_edit = gr.Image(
                            label="✨ Edited Result",
                            show_label=True,
                            interactive=False,
                            height=400,
                        )

                with gr.Row():
                    prompt_edit = gr.Textbox(
                        label="Editing Instruction",
                        placeholder="Describe how to edit the image (e.g., 'Add sunglasses', 'Change background to a sunset beach', 'Make it look like a painting')...",
                        lines=2,
                        max_lines=4,
                        scale=4,
                    )
                    run_edit = gr.Button("✨ Edit", variant="primary", scale=1)

                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    with gr.Row():
                        seed_edit = gr.Slider(
                            label="Seed",
                            minimum=0,
                            maximum=MAX_SEED,
                            step=1,
                            value=0,
                        )
                        randomize_seed_edit = gr.Checkbox(
                            label="Randomize seed", value=True
                        )

                    guidance_scale_edit = gr.Slider(
                        label="Guidance Scale",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.1,
                        value=2.5,
                        info="Higher = stronger adherence to prompt, lower = more creative",
                    )

                    num_steps_edit = gr.Slider(
                        label="Inference Steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=28,
                        info="More steps = better quality (28 is recommended for Kontext)",
                    )

                # Reuse edited image as input
                reuse_button = gr.Button("🔄 Use Result as New Input", visible=False)

                gr.Examples(
                    examples=[
                        [
                            "A beautiful landscape with mountains and a lake",
                            "Turn it into a winter scene with snow",
                        ],
                        [
                            "A portrait of a person",
                            "Add vintage film grain and warm color grading",
                        ],
                        [
                            "A city street",
                            "Transform into a cyberpunk neon-lit futuristic city at night",
                        ],
                        [
                            "A simple sketch of a house",
                            "Render it as a photorealistic modern architectural masterpiece at golden hour",
                        ],
                    ],
                    inputs=[input_image_edit, prompt_edit],
                    label="✨ Example Edits — Upload a matching image and try",
                )

    # ─── Footer ────────────────────────────────────────────────────────────
    gr.Markdown(
        """
        <div class="footer-info">
        <strong>The-X</strong> — Powered by FLUX.1-schnell & FLUX.1-Kontext-dev | 
        <a href="https://huggingface.co/black-forest-labs/FLUX.1-schnell" target="_blank">FLUX.1-schnell</a> (Apache 2.0) | 
        <a href="https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev" target="_blank">FLUX.1-Kontext-dev</a> (Non-commercial)
        </div>
        """
    )

    # ─── Event Handlers ────────────────────────────────────────────────────

    # Text-to-Image
    gr.on(
        triggers=[run_t2i.click, prompt_t2i.submit],
        fn=generate_image,
        inputs=[
            prompt_t2i,
            seed_t2i,
            randomize_seed_t2i,
            width_t2i,
            height_t2i,
            num_steps_t2i,
        ],
        outputs=[result_t2i, seed_t2i],
    ).then(
        fn=lambda: gr.Button(visible=True),
        outputs=[send_to_edit],
    )

    # Send generated image to edit tab
    send_to_edit.click(
        fn=lambda img: (img, gr.Tabs(selected="edit")),
        inputs=[result_t2i],
        outputs=[input_image_edit, tabs],
    )

    # Image Editing
    gr.on(
        triggers=[run_edit.click, prompt_edit.submit],
        fn=edit_image,
        inputs=[
            input_image_edit,
            prompt_edit,
            seed_edit,
            randomize_seed_edit,
            guidance_scale_edit,
            num_steps_edit,
        ],
        outputs=[result_edit, seed_edit],
    ).then(
        fn=lambda: gr.Button(visible=True),
        outputs=[reuse_button],
    )

    # Reuse edited image
    reuse_button.click(
        fn=lambda image: image,
        inputs=[result_edit],
        outputs=[input_image_edit],
    )


demo.launch(mcp_server=True)