import random

import gradio as gr
import spaces
import torch

from PIL import Image
from diffusers import ZImageImg2ImgPipeline


# ============================================================
# CONFIG
# ============================================================

MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"

# IMPORTANT:
# Z-Image strength=1.0 essentially ignores the input image.
# We intentionally cap the UI at 0.85 for better editing behavior.
MIN_STRENGTH = 0.15
MAX_STRENGTH = 0.85
DEFAULT_STRENGTH = 0.55

NUM_INFERENCE_STEPS = 8
GUIDANCE_SCALE = 0.0

MAX_IMAGE_SIZE = 1024

GPU_DURATION = 180


# ============================================================
# MODEL
# ============================================================

print(
    "Loading Z-Image-Turbo...",
    flush=True,
)

pipe = ZImageImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
)

pipe.to("cuda")

pipe.set_progress_bar_config(
    disable=True
)

print(
    "Z-Image-Turbo loaded successfully.",
    flush=True,
)


# ============================================================
# IMAGE PREPARATION
# ============================================================

def prepare_image(
    image: Image.Image,
) -> Image.Image:

    if image is None:
        raise ValueError(
            "Please upload an image."
        )

    image = image.convert("RGB")

    width, height = image.size

    if width < 64 or height < 64:
        raise ValueError(
            "Image must be at least 64×64 pixels."
        )

    # Keep the original aspect ratio.
    scale = min(
        MAX_IMAGE_SIZE / width,
        MAX_IMAGE_SIZE / height,
        1.0,
    )

    new_width = int(width * scale)
    new_height = int(height * scale)

    # Z-Image works with dimensions divisible by 8.
    new_width = max(
        64,
        (new_width // 8) * 8,
    )

    new_height = max(
        64,
        (new_height // 8) * 8,
    )

    if (
        new_width != image.width
        or new_height != image.height
    ):
        image = image.resize(
            (new_width, new_height),
            Image.Resampling.LANCZOS,
        )

    return image


# ============================================================
# PROMPT HANDLING
# ============================================================

def build_edit_prompt(
    prompt: str,
    strength: float,
) -> str:

    prompt = prompt.strip()

    # Lower strength = stronger preservation.
    if strength <= 0.35:

        prefix = (
            "Edit the provided image while preserving "
            "the original subject, identity, composition, "
            "camera angle, pose, framing, and overall "
            "structure. Make only the requested visual "
            "changes. "
        )

    elif strength <= 0.60:

        prefix = (
            "Transform the provided image according to "
            "the request while preserving the main subject, "
            "recognizable identity, composition, and "
            "important visual structure. "
        )

    else:

        prefix = (
            "Creatively transform the provided image "
            "according to the request while keeping the "
            "main subject and important visual identity "
            "recognizable. Do not completely replace the "
            "original concept unless explicitly requested. "
        )

    return prefix + prompt


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU(
    duration=GPU_DURATION
)
def generate(
    image,
    prompt,
    strength,
    seed,
):

    if image is None:
        raise gr.Error(
            "Please upload an image first."
        )

    if not prompt or not prompt.strip():
        raise gr.Error(
            "Please enter a transformation prompt."
        )

    # --------------------------------------------------------
    # Strength
    # --------------------------------------------------------

    try:
        strength = float(strength)
    except Exception:

        raise gr.Error(
            "Invalid transformation strength."
        )

    strength = max(
        MIN_STRENGTH,
        min(
            MAX_STRENGTH,
            strength,
        ),
    )

    # --------------------------------------------------------
    # Seed
    # --------------------------------------------------------

    try:
        seed = int(seed)
    except Exception:

        seed = -1

    if seed < 0:

        seed = random.SystemRandom().randint(
            0,
            2**32 - 1,
        )

    # --------------------------------------------------------
    # Generation
    # --------------------------------------------------------

    try:

        image = prepare_image(image)

        final_prompt = build_edit_prompt(
            prompt,
            strength,
        )

        generator = torch.Generator(
            device="cuda"
        ).manual_seed(seed)

        print(
            "Generation started | "
            f"seed={seed} | "
            f"strength={strength}",
            flush=True,
        )

        print(
            f"Prompt: {final_prompt}",
            flush=True,
        )

        with torch.inference_mode():

            result = pipe(
                prompt=final_prompt,
                image=image,
                strength=strength,
                num_inference_steps=NUM_INFERENCE_STEPS,
                guidance_scale=GUIDANCE_SCALE,
                generator=generator,
            )

        output = result.images[0]

        print(
            "Generation complete | "
            f"seed={seed}",
            flush=True,
        )

        return output, seed

    except torch.cuda.OutOfMemoryError:

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        raise gr.Error(
            "GPU memory was insufficient. "
            "Try a smaller input image."
        )

    except Exception as exc:

        print(
            "Generation error:",
            repr(exc),
            flush=True,
        )

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        raise gr.Error(
            f"Generation failed: {exc}"
        )


# ============================================================
# UI
# ============================================================

CSS = """
#title {
    text-align: center;
}

#subtitle {
    text-align: center;
    opacity: 0.7;
}

.generate-btn {
    min-height: 52px;
    font-size: 18px !important;
    font-weight: 600 !important;
}

.strength-info {
    text-align: center;
    opacity: 0.7;
    font-size: 13px;
}
"""


with gr.Blocks(
    title="IMAGINE — Z-Image Image to Image",
    css=CSS,
) as demo:

    gr.Markdown(
        "# IMAGINE",
        elem_id="title",
    )

    gr.Markdown(
        "Z-Image-Turbo Image Transformation",
        elem_id="subtitle",
    )

    with gr.Row():

        # ====================================================
        # INPUT
        # ====================================================

        with gr.Column():

            input_image = gr.Image(
                type="pil",
                label="Upload Image",
            )

            prompt = gr.Textbox(
                label="Transformation Prompt",
                placeholder=(
                    "Example: Change the clothes to "
                    "a black formal suit while keeping "
                    "the person's face and pose."
                ),
                lines=4,
            )

            strength = gr.Slider(
                minimum=MIN_STRENGTH,
                maximum=MAX_STRENGTH,
                value=DEFAULT_STRENGTH,
                step=0.05,
                label="Transformation Strength",
                info=(
                    "Lower = preserve more • "
                    "Higher = transform more"
                ),
            )

            gr.Markdown(
                "Preserve  ←────────  Balanced  ────────→  Transform",
                elem_classes="strength-info",
            )

            seed = gr.Number(
                value=-1,
                precision=0,
                label="Seed",
                info="-1 = random",
            )

            generate_button = gr.Button(
                "Generate Image",
                variant="primary",
                elem_classes="generate-btn",
            )

        # ====================================================
        # OUTPUT
        # ====================================================

        with gr.Column():

            output_image = gr.Image(
                type="pil",
                label="Generated Image",
            )

            output_seed = gr.Number(
                label="Used Seed",
                interactive=False,
            )


    # ========================================================
    # EXAMPLES
    # ========================================================

    gr.Examples(
        examples=[
            [
                (
                    "Change the lighting to warm golden "
                    "hour lighting while keeping the "
                    "same person, face, pose, background "
                    "and composition."
                ),
                0.30,
                -1,
            ],
            [
                (
                    "Change the person's clothing to a "
                    "premium black formal outfit while "
                    "keeping the same face, pose and "
                    "composition."
                ),
                0.50,
                -1,
            ],
            [
                (
                    "Transform this scene into a cinematic "
                    "fantasy environment while keeping "
                    "the main subject recognizable."
                ),
                0.70,
                -1,
            ],
        ],
        inputs=[
            prompt,
            strength,
            seed,
        ],
    )


    # ========================================================
    # API
    # ========================================================

    generate_button.click(
        fn=generate,
        inputs=[
            input_image,
            prompt,
            strength,
            seed,
        ],
        outputs=[
            output_image,
            output_seed,
        ],
        api_name="generate",
    )


# ============================================================
# LAUNCH
# ============================================================

if __name__ == "__main__":

    demo.queue(
        max_size=20,
        default_concurrency_limit=1,
    ).launch(
        server_name="0.0.0.0",
        server_port=7860,
    )