import os
import uuid
import numpy as np
import random
import tempfile
import zipfile

import spaces
import torch
import gradio as gr

from PIL import Image
from pptx import Presentation
from huggingface_hub import login

from diffusers import QwenImageLayeredPipeline
from diffusers.quantizers import PipelineQuantizationConfig


# ============================================================
# Configuration
# ============================================================

LOG_DIR = "/tmp/local"
MAX_SEED = np.iinfo(np.int32).max

MODEL_ID = "Qwen/Qwen-Image-Layered"

# Compute dtype used by the quantized model.
DTYPE = torch.bfloat16


# ============================================================
# Hugging Face authentication
# ============================================================

HF_TOKEN = os.environ.get("hf")

if HF_TOKEN:
    login(token=HF_TOKEN)


# ============================================================
# 4-BIT QUANTIZATION
# ============================================================
#
# The original Qwen-Image-Layered model is too large to fit
# comfortably into the free ZeroGPU 48 GB allocation.
#
# We quantize:
#
#   transformer   -> 4-bit NF4
#   text_encoder  -> 4-bit NF4
#
# Computation remains BF16.
#
# ============================================================

quant_config = PipelineQuantizationConfig(
    quant_backend="bitsandbytes_4bit",
    quant_kwargs={
        "load_in_4bit": True,
        "bnb_4bit_quant_type": "nf4",
        "bnb_4bit_compute_dtype": DTYPE,
        "bnb_4bit_use_double_quant": True,
    },
    components_to_quantize=[
        "transformer",
        "text_encoder",
    ],
)


# ============================================================
# Load pipeline
# ============================================================
#
# IMPORTANT:
# This remains at module level.
# ZeroGPU specifically optimizes CUDA placement performed
# during Space startup.
#
# device_map="cuda" ensures the quantized pipeline is placed
# on CUDA without first loading the complete BF16 model onto
# the GPU.
# ============================================================

print("Loading Qwen-Image-Layered with 4-bit quantization...")

pipeline = QwenImageLayeredPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE,
    quantization_config=quant_config,
    device_map="cuda",
)

print("Qwen-Image-Layered loaded successfully.")


# ============================================================
# Utility functions
# ============================================================

def ensure_dirname(path: str):
    if path and not os.path.exists(path):
        os.makedirs(path, exist_ok=True)


def random_str(length=8):
    return uuid.uuid4().hex[:length]


def get_image_from_input(input_image):
    """
    Convert the Gradio image input into a PIL RGBA image.
    """

    if isinstance(input_image, list):
        if not input_image:
            raise ValueError("No input image provided.")
        input_image = input_image[0]

    if isinstance(input_image, str):
        return Image.open(input_image).convert("RGB").convert("RGBA")

    if isinstance(input_image, Image.Image):
        return input_image.convert("RGB").convert("RGBA")

    if isinstance(input_image, np.ndarray):
        return Image.fromarray(input_image).convert("RGB").convert("RGBA")

    raise ValueError(
        f"Unsupported input_image type: {type(input_image)}"
    )


# ============================================================
# PPTX export
# ============================================================

def imagelist_to_pptx(img_files):
    with Image.open(img_files[0]) as img:
        img_width_px, img_height_px = img.size

    def px_to_emu(px, dpi=96):
        inch = px / dpi
        emu = inch * 914400
        return int(emu)

    prs = Presentation()

    prs.slide_width = px_to_emu(img_width_px)
    prs.slide_height = px_to_emu(img_height_px)

    slide = prs.slides.add_slide(prs.slide_layouts[6])

    left = top = 0

    for img_path in img_files:
        slide.shapes.add_picture(
            img_path,
            left,
            top,
            width=px_to_emu(img_width_px),
            height=px_to_emu(img_height_px),
        )

    with tempfile.NamedTemporaryFile(
        suffix=".pptx",
        delete=False
    ) as tmp:
        prs.save(tmp.name)
        return tmp.name


# ============================================================
# ZIP export
# ============================================================

def export_gallery_zip(images):
    """
    Create a ZIP containing all generated layers.
    """

    images = [e[0] for e in images]

    with tempfile.NamedTemporaryFile(
        suffix=".zip",
        delete=False
    ) as tmp:

        with zipfile.ZipFile(
            tmp.name,
            "w",
            zipfile.ZIP_DEFLATED
        ) as zipf:

            for i, img_path in enumerate(images):

                ext = (
                    os.path.splitext(img_path)[1]
                    or ".png"
                )

                zipf.write(
                    img_path,
                    f"layer_{i + 1}{ext}"
                )

        return tmp.name


# ============================================================
# Main inference function
# ============================================================

@spaces.GPU(duration=180)
def infer(
    input_image,
    seed: int = 777,
    randomize_seed: bool = False,
    prompt: str | None = None,
    neg_prompt: str = " ",
    true_guidance_scale: float = 4.0,
    num_inference_steps: int = 50,
    layer: int = 4,
    cfg_norm: bool = True,
    use_en_prompt: bool = True,
):
    """
    Decompose an input image into multiple editable layers.

    Args:
        input_image: Input image.
        seed: Random seed.
        randomize_seed: Generate a random seed when enabled.
        prompt: Optional description of the input image.
        neg_prompt: Negative prompt.
        true_guidance_scale: CFG guidance scale.
        num_inference_steps: Number of diffusion steps.
        layer: Number of output layers.
        cfg_norm: Enable CFG normalization.
        use_en_prompt: Automatically use English prompting.

    Returns:
        Generated layer images, PPTX file, and ZIP file.
    """

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

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

    # --------------------------------------------------------
    # Convert input image
    # --------------------------------------------------------

    pil_image = get_image_from_input(input_image)

    # --------------------------------------------------------
    # Pipeline inputs
    # --------------------------------------------------------

    inputs = {
        "image": pil_image,

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

        "true_cfg_scale": true_guidance_scale,

        "prompt": prompt,

        "negative_prompt": neg_prompt,

        "num_inference_steps": int(
            num_inference_steps
        ),

        "num_images_per_prompt": 1,

        "layers": int(layer),

        # Qwen-Image-Layered recommended resolution.
        "resolution": 640,

        "cfg_normalize": cfg_norm,

        "use_en_prompt": use_en_prompt,
    }

    print(
        f"Running inference | "
        f"seed={seed} | "
        f"steps={num_inference_steps} | "
        f"layers={layer}"
    )

    # --------------------------------------------------------
    # GPU inference
    # --------------------------------------------------------

    with torch.inference_mode():

        result = pipeline(**inputs)

        output_images = result.images[0]

    # --------------------------------------------------------
    # Save generated layers
    # --------------------------------------------------------

    output = []
    temp_files = []

    for image in output_images:

        output.append(image)

        tmp = tempfile.NamedTemporaryFile(
            suffix=".png",
            delete=False
        )

        image.save(tmp.name)

        temp_files.append(tmp.name)

        tmp.close()

    # --------------------------------------------------------
    # Create PPTX
    # --------------------------------------------------------

    pptx_path = imagelist_to_pptx(
        temp_files
    )

    # --------------------------------------------------------
    # Create ZIP
    # --------------------------------------------------------

    zip_path = export_gallery_zip(
        [
            (path, f"Layer {i + 1}")
            for i, path in enumerate(temp_files)
        ]
    )

    return (
        output,
        pptx_path,
        zip_path,
    )


# ============================================================
# Setup
# ============================================================

ensure_dirname(LOG_DIR)


# ============================================================
# Example images
# ============================================================

examples = [
    "assets/test_images/1.png",
    "assets/test_images/2.png",
    "assets/test_images/3.png",
    "assets/test_images/4.png",
    "assets/test_images/5.png",
    "assets/test_images/6.png",
    "assets/test_images/7.png",
    "assets/test_images/8.png",
    "assets/test_images/9.png",
    "assets/test_images/10.png",
    "assets/test_images/11.png",
    "assets/test_images/12.png",
    "assets/test_images/13.png",
]


# ============================================================
# Gradio UI
# ============================================================

with gr.Blocks() as demo:

    with gr.Column(elem_id="col-container"):

        gr.HTML(
            '''
            <img
                src="https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-Image/layered/qwen-image-layered-logo.png"
                alt="Qwen-Image-Layered Logo"
                width="600"
                style="display:block;margin:0 auto;"
            >
            '''
        )

        gr.Markdown(
            """
            The text prompt is intended to describe the overall
            content of the input image—including elements that
            may be partially occluded.

            It is not designed to control the semantic content
            of individual layers explicitly.
            """
        )

        with gr.Row():

            # ------------------------------------------------
            # INPUT COLUMN
            # ------------------------------------------------

            with gr.Column(scale=1):

                input_image = gr.Image(
                    label="Input Image",
                    image_mode="RGBA",
                    type="filepath",
                )

                with gr.Accordion(
                    "Advanced Settings",
                    open=False
                ):

                    prompt = gr.Textbox(
                        label="Prompt (Optional)",
                        placeholder=(
                            "Please enter the prompt "
                            "to describe the image."
                        ),
                        value="",
                        lines=2,
                    )

                    neg_prompt = gr.Textbox(
                        label="Negative Prompt (Optional)",
                        placeholder=(
                            "Please enter the negative prompt"
                        ),
                        value=" ",
                        lines=2,
                    )

                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=0,
                    )

                    randomize_seed = gr.Checkbox(
                        label="Randomize seed",
                        value=True,
                    )

                    true_guidance_scale = gr.Slider(
                        label="True guidance scale",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.1,
                        value=4.0,
                    )

                    num_inference_steps = gr.Slider(
                        label="Number of inference steps",
                        minimum=1,
                        maximum=50,
                        step=1,
                        value=50,
                    )

                    layer = gr.Slider(
                        label="Layers",
                        minimum=2,
                        maximum=10,
                        step=1,
                        value=4,
                    )

                    cfg_norm = gr.Checkbox(
                        label="Whether enable CFG normalization",
                        value=True,
                    )

                    use_en_prompt = gr.Checkbox(
                        label=(
                            "Automatic caption language if no "
                            "prompt is provided, True for EN, "
                            "False for ZH"
                        ),
                        value=True,
                    )

                run_button = gr.Button(
                    "Decompose!",
                    variant="primary",
                )

            # ------------------------------------------------
            # OUTPUT COLUMN
            # ------------------------------------------------

            with gr.Column(scale=2):

                gallery = gr.Gallery(
                    label="Layers",
                    columns=4,
                    rows=1,
                    format="png",
                )

                with gr.Row():

                    export_file = gr.File(
                        label="Download PPTX"
                    )

                    export_zip_file = gr.File(
                        label="Download ZIP"
                    )

    # ========================================================
    # Examples
    # ========================================================

    gr.Examples(
        examples=examples,
        inputs=[input_image],
        outputs=[
            gallery,
            export_file,
            export_zip_file,
        ],
        fn=infer,
        examples_per_page=14,
        cache_examples=False,
        run_on_click=True,
    )

    # ========================================================
    # API ENDPOINT
    # ========================================================

    run_button.click(
        fn=infer,
        inputs=[
            input_image,
            seed,
            randomize_seed,
            prompt,
            neg_prompt,
            true_guidance_scale,
            num_inference_steps,
            layer,
            cfg_norm,
            use_en_prompt,
        ],
        outputs=[
            gallery,
            export_file,
            export_zip_file,
        ],
        api_name="infer",
    )


# ============================================================
# Launch
# ============================================================

if __name__ == "__main__":

    demo.launch(
        mcp_server=True,
    )