import gradio as gr
import math
import numpy as np
import random
import torch
import spaces
import os
import requests
import tempfile
import shutil

from PIL import Image
from diffusers import QwenImageEditPlusPipeline
from urllib.parse import urlparse
from pathlib import Path

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

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511",
    torch_dtype=dtype
).to(device)

pipe.load_lora_weights(
    "lightx2v/Qwen-Image-Edit-2511-Lightning",
    weight_name="Qwen-Image-Edit-2511-Lightning-4steps-V1.0-bf16.safetensors",
)
pipe.fuse_lora()
pipe.unload_lora_weights()

_VAE_IMAGE_SIZE = 1024 * 1024


def calculate_vae_gen_size(image: Image.Image) -> tuple:
    W, H = image.size
    ratio = W / H
    gen_w = math.sqrt(_VAE_IMAGE_SIZE * ratio)
    gen_h = gen_w / ratio
    gen_w = round(gen_w / 32) * 32
    gen_h = round(gen_h / 32) * 32
    return int(gen_w), int(gen_h)


def resize_image(image: Image.Image) -> Image.Image:
    MAX_SIDE = 1328
    w, h = image.size
    scale = min(MAX_SIDE / w, MAX_SIDE / h, 1.0)
    new_w = (int(w * scale) // 16) * 16
    new_h = (int(h * scale) // 16) * 16
    if (new_w, new_h) == (w, h):
        return image
    return image.resize((new_w, new_h), Image.LANCZOS)


def load_lora_auto(pipe, lora_input: str):
    lora_input = lora_input.strip()
    if not lora_input:
        return False

    if "/" in lora_input and not lora_input.startswith("http"):
        pipe.load_lora_weights(lora_input)
        return True

    if lora_input.startswith("http"):
        url = lora_input
        if "huggingface.co" in url and "/blob/" not in url and "/resolve/" not in url:
            repo_id = urlparse(url).path.strip("/")
            pipe.load_lora_weights(repo_id)
            return True
        if "/blob/" in url:
            url = url.replace("/blob/", "/resolve/")
        tmp_dir = tempfile.mkdtemp()
        local_path = os.path.join(tmp_dir, os.path.basename(urlparse(url).path))
        try:
            resp = requests.get(url, stream=True)
            resp.raise_for_status()
            with open(local_path, "wb") as f:
                for chunk in resp.iter_content(chunk_size=8192):
                    f.write(chunk)
            pipe.load_lora_weights(local_path)
            return True
        finally:
            shutil.rmtree(tmp_dir, ignore_errors=True)

    return False


def save_outputs(images, seed, out_format):
    out_dir = Path("output")
    out_dir.mkdir(parents=True, exist_ok=True)
    paths = []
    ext = "jpg" if out_format == "jpg" else "png"
    pil_format = "JPEG" if out_format == "jpg" else "PNG"

    for i, img in enumerate(images[:2]):
        path = out_dir / f"gen_{seed}_{i}.{ext}"
        if pil_format == "JPEG":
            img.convert("RGB").save(path, format="JPEG", quality=100, subsampling=0, optimize=False)
        else:
            img.convert("RGB").save(path, format="PNG")
        paths.append(str(path))

    return paths


@spaces.GPU
def infer(
    gallery_images,
    prompt: str,
    lora_id: str = "",
    seed: int = 0,
    randomize_seed: bool = True,
    true_guidance_scale: float = 1.0,
    num_inference_steps: int = 4,
    width: int = 1024,
    height: int = 1024,
    auto_size: bool = True,
    output_format: str = "png",
    progress=gr.Progress(track_tqdm=True)
):
    if not gallery_images:
        raise gr.Error("Please upload at least 1 image.")

    processed_images = []
    for item in gallery_images[:3]:
        img_obj = item[0] if isinstance(item, tuple) else (item.image if hasattr(item, 'image') else item)
        processed_images.append(resize_image(img_obj).convert("RGB"))

    images = processed_images

    if len(gallery_images) > 3:
        gr.Warning("Only the first 3 images are used.")

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

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

    custom_lora_loaded = False
    if lora_id and lora_id.strip():
        try:
            custom_lora_loaded = load_lora_auto(pipe, lora_id)
        except Exception as e:
            print(f"LoRA load failed: {e}")
            custom_lora_loaded = False

    if auto_size:
        width, height = calculate_vae_gen_size(images[0])

    try:
        result = pipe(
            image=images,
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            generator=generator,
            true_cfg_scale=true_guidance_scale,
            num_images_per_prompt=2,
        ).images

        out_paths = save_outputs(result, seed, output_format)

    finally:
        if custom_lora_loaded:
            pipe.unload_lora_weights()

    return out_paths, seed


css = '''
.gradio-container,
.gradio-container * {
    font-family: Helvetica, Arial, sans-serif !important;
}
#col-container { max-width: 1000px; margin: 0 auto; }
#examples { max-width: 1000px; margin: 0 auto; }
.image-container { min-height: 300px; }
.dark,
.dark * {
    color-scheme: dark;
}
.dark .progress-text,
.dark label,
.dark .label,
.dark .wrap label,
.dark .gradio-container label,
.dark .prose,
.dark .prose *,
.dark .markdown,
.dark .markdown *,
.dark p,
.dark span,
.dark div,
.dark small,
.dark strong,
.dark em,
.dark h1,
.dark h2,
.dark h3,
.dark h4,
.dark h5,
.dark h6 {
    color: #808080 !important;
}
.dark input,
.dark textarea,
.dark select,
.dark .input textarea,
.dark .input input {
    color: #808080 !important;
    caret-color: #808080 !important;
}
.dark a,
.dark a * {
    color: #808080 !important;
}
.dark svg,
.dark svg * {
    fill: #808080 !important;
    stroke: #808080 !important;
}
#quick-loras-container {
    display: flex !important;
    flex-wrap: wrap !important;
    align-items: center !important;
    gap: 6px !important;
    padding: 4px 0 !important;
}
#quick-loras-container > .form,
#quick-loras-container > div {
    flex: 0 0 auto !important;
    width: auto !important;
    min-width: 0 !important;
    padding: 0 !important;
    background: none !important;
    border: none !important;
    box-shadow: none !important;
    gap: 0 !important;
}
#quick-loras-container button {
    width: auto !important;
    min-width: fit-content !important;
    white-space: nowrap !important;
    background: rgba(255,255,255,0.06) !important;
    border: 1px solid rgba(255,255,255,0.12) !important;
    color: #808080 !important;
    box-shadow: none !important;
}
#quick-loras-container button:hover {
    background: rgba(255,255,255,0.10) !important;
    border-color: rgba(255,255,255,0.20) !important;
}
.quick-lora-link {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    width: 28px;
    height: 28px;
    border-radius: 6px;
    background: rgba(255,255,255,0.08);
    color: #808080 !important;
    text-decoration: none;
    font-size: 14px;
    line-height: 1;
    transition: background 0.15s ease, transform 0.1s ease;
    flex-shrink: 0;
    vertical-align: middle;
}
.quick-lora-link:hover {
    background: rgba(255,255,255,0.16);
    transform: scale(1.12);
}
.dark .gr-button,
.dark button {
    color: #808080 !important;
}
.dark .gr-button-primary,
.dark .gr-button-secondary {
    color: #808080 !important;
}
'''

POPULAR_LORAS = [
    (
        "1",
        "ovi054/QIE-2511-Color-Grade-Transfer-LoRA",
        "Transfer ONLY the color grading from Image 2 onto Image 1",
        "https://huggingface.co/ovi054/QIE-2511-Color-Grade-Transfer-LoRA",
    ),
    (
        "2",
        "fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA",
        "<sks> front-left quarter view elevated shot medium shot",
        "https://huggingface.co/fal/Qwen-Image-Edit-2511-Multiple-Angles-LoRA",
    ),
    (
        "3",
        "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap/resolve/main/bfs_head_v5_2511_original.safetensors",
        "face swap face from Image 1 to Image 2.",
        "https://huggingface.co/Alissonerdx/BFS-Best-Face-Swap",
    ),
    (
        "4",
        "https://www.modelscope.ai/models/Playmaker/floatfit3d/resolve/master/floatfit3d_25.safetensors",
        "extract the outfit from the person and render it as a floating 3d clothing display on a gray background.",
        "https://huggingface.co/Playmaker/floatfit3d",
    ),
    (
        "5",
        "https://huggingface.co/prithivMLmods/Qwen-Image-Edit-2511-Unblur-Upscale/resolve/main/Qwen-Image-Edit-Unblur-Upscale_20.safetensors",
        "unblur and upscale",
        "https://huggingface.co/prithivMLmods/Qwen-Image-Edit-2511-Unblur-Upscale",
    ),
    (
        "6",
        "dx8152/Qwen-Image-Edit-2511-Style-Transfer",
        "Change the style of Figure 1 to the style of Figure 2.",
        "https://huggingface.co/dx8152/Qwen-Image-Edit-2511-Style-Transfer",
    ),
    (
        "7",
        "https://huggingface.co/ilkerzgi/krea-2-emerald-noir-oil-lora.safetensors",
        "Emerald noir oil style.",
        "https://huggingface.co/ilkerzgi/krea-2-emerald-noir-oil-lora",
    ),
    (
        "8",
        "lilylilith/AnyPose",
        "Make the person in image 1 do the exact same pose of the person in image 2. Changing the style and background of the image of the person in image 1 is undesirable, so don't do it. The new pose should be pixel accurate to the pose we are trying to copy. The position of the arms and head and legs should be the same as the pose we are trying to copy. Change the field of view and angle to match exactly image 2. Head tilt and eye gaze pose should match the person in image 2.",
        "https://huggingface.co/lilylilith/AnyPose",
    ),
]

neutral_dark_theme = gr.themes.Monochrome(
    primary_hue=gr.themes.colors.neutral,
    secondary_hue=gr.themes.colors.neutral,
    neutral_hue=gr.themes.colors.neutral,
).set(
    body_background_fill="#000000",
    body_background_fill_dark="#000000",
    background_fill_primary="#000000",
    background_fill_primary_dark="#000000",
    background_fill_secondary="#000000",
    background_fill_secondary_dark="#000000",
    block_background_fill="#000000",
    block_background_fill_dark="#000000",
    input_background_fill="#000000",
    input_background_fill_dark="#000000",
    block_border_color="#555555",
    block_border_color_dark="#555555",
    input_border_color="#555555",
    input_border_color_dark="#555555",
    body_text_color="#555555",
    body_text_color_dark="#555555",
    body_text_color_subdued="#555555",
    body_text_color_subdued_dark="#555555",
    block_label_text_color="#555555",
    block_label_text_color_dark="#555555",
    block_title_text_color="#555555",
    block_title_text_color_dark="#555555",
    button_primary_background_fill="#000000",
    button_primary_background_fill_dark="#000000",
    button_primary_background_fill_hover="#000000",
    button_primary_background_fill_hover_dark="#000000",
    button_primary_text_color="#555555",
    button_primary_text_color_dark="#555555",
    button_secondary_background_fill="#000000",
    button_secondary_background_fill_dark="#000000",
    button_secondary_background_fill_hover="#000000",
    button_secondary_background_fill_hover_dark="#000000",
    button_secondary_text_color="#555555",
    button_secondary_text_color_dark="#555555",
    link_text_color="#555555",
    link_text_color_dark="#555555",
    table_text_color="#555555",
    table_text_color_dark="#555555",
    slider_color="#555555",
    slider_color_dark="#555555",
)

with gr.Blocks(theme=neutral_dark_theme, css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("")
        gr.Markdown(
            ""
        )

        with gr.Row():
            with gr.Column():
                input_gallery = gr.Gallery(
                    label="",
                    columns=3,
                    rows=1,
                    object_fit="contain",
                    type="pil",
                    interactive=True,
                )

                prompt = gr.Textbox(label="", lines=2)

                lora_id = gr.Textbox(
                    label="",
                    info="",
                    placeholder="",
                )

                output_format = gr.Checkbox(label="", value=False)

                run_btn = gr.Button("Edit", variant="primary", size="lg")

                with gr.Accordion("", open=False):
                    auto_size = gr.Checkbox(label="", value=True)
                    with gr.Row():
                        width = gr.Slider(label="", value=2048, minimum=1024, maximum=2048, step=256)
                        height = gr.Slider(label="", value=2048, minimum=1024, maximum=2048, step=256)
                    seed = gr.Slider(label="", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="", value=True)
                    true_guidance_scale = gr.Slider(
                        label="", minimum=1, maximum=10, step=1, value=1
                    )
                    num_inference_steps = gr.Slider(
                        label="", minimum=1, maximum=40, step=1, value=4
                    )

            with gr.Column():
                result = gr.Gallery(label="", columns=4, rows=4, object_fit="contain")
                gr.Markdown("")
                with gr.Row(elem_id="quick-loras-container"):
                    for btn_label, repo, trigger, url in POPULAR_LORAS:
                        gr.Button(btn_label, size="sm", variant="secondary").click(
                            fn=lambda r=repo, t=trigger: (r, t),
                            outputs=[lora_id, prompt]
                        )
                        
    run_btn.click(
        fn=infer,
        inputs=[input_gallery, prompt, lora_id, seed, randomize_seed, true_guidance_scale, num_inference_steps, width, height, auto_size, output_format],
        outputs=[result, seed]
    )

demo.launch(mcp_server=False, theme=neutral_dark_theme, css=css)