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 typing import List, Tuple
from urllib.parse import urlparse

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

# --- Model Loading ---
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)

# Fuse the lightning LoRA directly into the base weights
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:
    """Cap longest side to 1328px, snap to multiples of 16."""
    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):
    """Load LoRA from HuggingFace repo ID, URL, or blob link."""
    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:
            print(f"Downloading LoRA from {url}...")
            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


@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,
    progress=gr.Progress(track_tqdm=True)
) -> Tuple[Image.Image, int]:
    # gallery_images is a list of (pil_image, caption) tuples or just pil images
    if not gallery_images:
        raise gr.Error("Please upload at least 1 image.")

    # images = [resize_image(img[0] if isinstance(img, tuple) else img).convert("RGB")
    #           for img in gallery_images[:3]]

    processed_images = []
    for item in gallery_images[:3]:
        # Gradio gallery yields (image, caption) tuples or dictionaries depending on version
        img_obj = item[0] if isinstance(item, tuple) else (item.image if hasattr(item, 'image') else item)
        
        # Apply your image scaling constraints and convert to RGB
        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)

    print(f"Running with {len(images)} input image(s), prompt: {prompt!r}")

    custom_lora_loaded = False
    if lora_id and lora_id.strip():
        try:
            custom_lora_loaded = load_lora_auto(pipe, lora_id)
            print(f"Loaded custom LoRA: {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=1,
        ).images[0]
    finally:
        if custom_lora_loaded:
            pipe.unload_lora_weights()

    return result, seed


# # --- UI ---
# css = "#col-container { max-width: 1100px; margin: 0 auto; }"


# --- UI ---
css = '''
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .progress-text { color: white !important }
#examples { max-width: 1000px; margin: 0 auto; }
.image-container { min-height: 300px; }

/* Quick LoRAs compact strip */
#quick-loras-container {
    display: flex !important;
    flex-wrap: wrap !important;
    align-items: center !important;
    gap: 6px !important;
    padding: 4px 0 !important;
}
/* Remove Gradio's default flex-stretch on each child wrapper */
#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;
}
/* Buttons stay content-width and use a neutral color (not yellow) */
#quick-loras-container button {
    width: auto !important;
    min-width: fit-content !important;
    white-space: nowrap !important;
    background: rgba(100,100,100,0.10) !important;
    border: 1px solid rgba(100,100,100,0.25) !important;
    color: inherit !important;
    box-shadow: none !important;
}
#quick-loras-container button:hover {
    background: rgba(100,100,100,0.18) !important;
    border-color: rgba(100,100,100,0.4) !important;
}
/* Keep LoRA names readable on the dark theme (default inherited color is too dark) */
.dark #quick-loras-container button {
    color: #ffffff !important;
}
/* Link icon pill */
.quick-lora-link {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    width: 28px;
    height: 28px;
    border-radius: 6px;
    background: rgba(128,128,128,0.12);
    color: inherit;
    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(128,128,128,0.28);
    transform: scale(1.12);
}


.space-grid {
    display: grid;
    grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
    gap: 16px;
    margin-top: 10px;
}

.space-card {
    display: flex;
    flex-direction: column;
    padding: 16px;
    border: 1px solid rgba(128,128,128,0.25);
    border-radius: 10px;
    text-decoration: none !important;
    color: inherit !important;
    background: rgba(128,128,128,0.05);
    transition: all 0.2s ease;
    cursor: pointer;
}

.space-card:hover {
    transform: translateY(-3px);
    border-color: #f97316; /* Matches the Citrus theme */
    box-shadow: 0 6px 12px rgba(0,0,0,0.08);
    background: rgba(128,128,128,0.1);
}

.space-title { 
    font-weight: 600; 
    font-size: 1.1em; 
    margin-bottom: 6px; 
}

.space-desc { 
    font-size: 0.9em; 
    opacity: 0.85; 
    line-height: 1.4; 
}

'''

# Each entry: (display_label, repo_id, trigger_words, model_url)
# model_url can be a HuggingFace, CivitAI, ModelScope, or any URL
POPULAR_LORAS = [
    (
        "🌓 Color Grade Transfer",
        "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",
    ),
    (
        "📐 Multiple Angles (Fal)",
        "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",
    ),
    (
        "🎭 BFS Best Face Swap",
        "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",
    ),
    (
        "👕 FloatFit3D (playmaker)",
        "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://www.modelscope.ai/models/Playmaker/floatfit3d",
    ),
    (
        "✨ Unblur Upscale",
        "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",
    ),
    (
        "🎨 Style-Transfer (dx8152)",
        "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",
    ),
    (
        "🔤 Letter LoRA",
        "https://www.modelscope.ai/models/krznun/LetterLora/resolve/master/LettersLora.safetensors",
        "extract lowercase font sheet for characters\na b c d e\nf g h i j\nk l m n o\np q r s t\nu v w x y\nz",
        "https://www.modelscope.ai/models/krznun/LetterLora",
    ),
    (
        "🧍 AnyPose",
        "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",
    ),
    (
        "💦 Gaussian Splash",
        "dx8152/Qwen-Image-Edit-2511-Gaussian-Splash",
        "高斯泼溅,参考图2的场景图，修复图1的场景图透视并修复空白区域",
        "https://huggingface.co/dx8152/Qwen-Image-Edit-2511-Gaussian-Splash",
    ),
    # (
    #     "➕ Object Adder",
    #     "prithivMLmods/Qwen-Image-Edit-2511-Object-Adder",
    #     "Add the specified objects to the image while preserving the background lighting and surrounding elements maintaining realism and original details.",
    #     "https://huggingface.co/prithivMLmods/Qwen-Image-Edit-2511-Object-Adder",
    # ),
    # (
    #     "📸 LumiPic HDR",
    #     "https://huggingface.co/oumoumad/LumiPic/resolve/main/hdrdit_v1_QE2511.safetensors",
    #     "Convert this image to HDR",
    #     "https://huggingface.co/oumoumad/LumiPic",
    # ),
    # (
    #     "💥 Torn Clothes",
    #     "nappa114514/Qwen-Image-Edit-2511-torn-clothes",
    #     "Tear the clothes according to the white areas of image2.",
    #     "https://huggingface.co/nappa114514/Qwen-Image-Edit-2511-torn-clothes",
    # ),
    # (
    #     "💡 Studio DeLight",
    #     "prithivMLmods/QIE-2511-Studio-DeLight",
    #     "Neutral uniform lighting Preserve identity and composition",
    #     "https://huggingface.co/prithivMLmods/QIE-2511-Studio-DeLight",
    # ),
    (
        "🖼️ Draw2Photo",
        "ovi054/QIE-2511-Draw2Photo-LoRA",
        "make it real",
        "https://huggingface.co/ovi054/QIE-2511-Draw2Photo-LoRA",
    ),
]

with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("## 🖌️ Qwen Image Edit 2511 + LoRA")
        gr.Markdown(
            "Upload **1–3 images** via the gallery, write a prompt, and optionally apply a LoRA. "
            "The order of images in the gallery is Image 1, 2, 3."
        )

        with gr.Row():
            with gr.Column():

                input_gallery = gr.Gallery(
                    label="Input Images (upload 1–3)",
                    columns=3,
                    rows=1,
                    # height=300,
                    object_fit="contain",
                    type="pil",
                    interactive=True,
                )

                prompt = gr.Textbox(
                    label="Prompt",
                    # placeholder="e.g. 'Put the person from Image 1 into the scene from Image 2'",
                    lines=2,
                )

                lora_id = gr.Textbox(label="Custom LoRA (optional)", info="URL or the path to the LoRA weights", placeholder="ovi054/QIE-2511-Color-Grade-Transfer-LoRA")

                

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

                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    auto_size = gr.Checkbox(label="Auto size", value=True)
                    with gr.Row():
                        width = gr.Slider(label="Width", value=1024, minimum=64, maximum=2048, step=16)
                        height = gr.Slider(label="Height", value=1024, minimum=64, maximum=2048, step=16)
                    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=1.0
                    )
                    num_inference_steps = gr.Slider(
                        label="Inference Steps", minimum=1, maximum=40, step=1, value=4
                    )

            with gr.Column():
                result = gr.Image(label="✨ Output", interactive=False)
                gr.Markdown("**Quick LoRAs:**")
                with gr.Row(elem_id="quick-loras-container"):
                    for btn_label, repo, trigger, url in POPULAR_LORAS:
                        # Button fills both lora_id and prompt with trigger words
                        gr.Button(btn_label, size="sm", variant="secondary").click(
                            fn=lambda r=repo, t=trigger: (r, t),
                            outputs=[lora_id, prompt]
                        )
                        # Link icon on the right — opens the model page in a new tab
                        gr.HTML(
                            f'<a href="{url}" target="_blank" rel="noopener noreferrer" '
                            f'class="quick-lora-link" title="Open model page">🔗</a>'
                        )
                # output_seed = gr.Number(label="Seed used", precision=0)

        gr.Markdown("---") 

        gr.HTML(
            """
            <h3 style='margin-bottom: 10px;'>🚀 Explore more of my spaces:</h3>
            <div class="space-grid">
                <a href="https://huggingface.co/spaces/build-small-hackathon/anim-vid-ai" target="_blank" class="space-card">
                    <div class="space-title">🎬 Anim Vid AI</div>
                    <div class="space-desc">Turn any topic into an engaging Manim animation video.</div>
                </a>
                <a href="https://huggingface.co/spaces/build-small-hackathon/Color-Grade-Transfer/" target="_blank" class="space-card">
                    <div class="space-title">🌗 Color Grade Transfer</div>
                    <div class="space-desc">Transfer Color Grade directly from a reference image to your source image.</div>
                </a>
            </div>
            """
        )
        

    run_btn.click(
        fn=infer,
        inputs=[input_gallery, prompt, lora_id, seed, randomize_seed, true_guidance_scale, num_inference_steps, width, height, auto_size],
        outputs=[result, seed]
    )

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=css)