import gradio as gr
import numpy as np
import random
import torch
import spaces
import os
from pathlib import Path

from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler
from optimization import optimize_pipeline_
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

import math

# --- Model Loading ---
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(3),
    "invert_sigmas": False,
    "max_image_seq_len": 8192,
    "max_shift": math.log(3),
    "num_train_timesteps": 1000,
    "shift": 1.0,
    "shift_terminal": None,
    "stochastic_sampling": False,
    "time_shift_type": "exponential",
    "use_beta_sigmas": False,
    "use_dynamic_shifting": True,
    "use_exponential_sigmas": False,
    "use_karras_sigmas": False,
}

scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", scheduler=scheduler, torch_dtype=dtype)

# Load the texture LoRA
pipe.load_lora_weights("2vXpSwA7/iroiro-lora",
        weight_name="qwen_lora/qie2509_lora_katame_transferring_01.safetensors", adapter_name="texture")
pipe.load_lora_weights("lightx2v/Qwen-Image-Lightning",
        weight_name="Qwen-Image-Lightning-4steps-V2.0-bf16.safetensors", adapter_name="lightning")
pipe.set_adapters(["texture", "lightning"], adapter_weights=[1.2, 1.])
pipe.fuse_lora(adapter_names=["texture", "lightning"], lora_scale=1)
pipe.unload_lora_weights()

pipe.transformer.__class__ = QwenImageTransformer2DModel
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

pipe.to(device)

optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")

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

# --- Load sample images ---
def get_sample_images(folder):
    """Get all image files from a folder."""
    folder_path = Path(folder)
    if not folder_path.exists():
        return []

    image_extensions = {'.png', '.jpg', '.jpeg', '.webp', '.bmp'}
    images = []
    for file in sorted(folder_path.iterdir()):
        if file.suffix.lower() in image_extensions:
            images.append(str(file))
    return images

slotA_images = get_sample_images("samples/slotA")
slotB_images = get_sample_images("samples/slotB")

def calculate_dimensions(image):
    """Calculate output dimensions based on content image, keeping largest side at 1024."""
    if image is None:
        return 1024, 1024
    
    original_width, original_height = image.size
    
    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
    
    # Ensure dimensions are multiples of 8
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return new_width, new_height

@spaces.GPU
def apply_texture(
    content_image,
    texture_image,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=False,
    num_inference_steps=4,
    progress=gr.Progress(track_tqdm=True)
):
    if content_image is None:
        raise gr.Error("Please upload a content image.")
    if texture_image is None:
        raise gr.Error("Please upload a texture image.")
    if not prompt or not prompt.strip():
        prompt = "change image1 character texture to image2 texture"
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # Calculate dimensions based on content image
    width, height = calculate_dimensions(content_image)
    
    # Prepare images
    content_pil = content_image.convert("RGB") if isinstance(content_image, Image.Image) else Image.open(content_image.name).convert("RGB")
    texture_pil = texture_image.convert("RGB") if isinstance(texture_image, Image.Image) else Image.open(texture_image.name).convert("RGB")
    
    pil_images = [content_pil, texture_pil]
    
    result = pipe(
        image=pil_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]
    
    return result, seed

# --- UI ---
css = '''
#col-container, #examples {
    max-width: 1400px;
    margin: 0 auto;
    padding: 20px;
}
.dark .progress-text{
    color: white !important;
}
/* Card style for image containers */
.image-card {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border-radius: 16px;
    padding: 4px;
    box-shadow: 0 8px 16px rgba(0,0,0,0.1);
}
/* Input section styling */
.input-section {
    background: rgba(255,255,255,0.05);
    border-radius: 12px;
    padding: 20px;
    margin-bottom: 15px;
}
/* Button styling */
.generate-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    font-size: 18px !important;
    font-weight: 600 !important;
    padding: 12px 24px !important;
    border-radius: 8px !important;
    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important;
    transition: all 0.3s ease !important;
}
.generate-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 6px 16px rgba(102, 126, 234, 0.6) !important;
}
/* Output section */
.output-section {
    background: rgba(255,255,255,0.03);
    border-radius: 12px;
    padding: 20px;
    min-height: 600px;
}
/* Accordion styling */
.accordion {
    border-radius: 8px;
    margin-top: 10px;
}
/* Image upload area */
.image-upload {
    border: 2px dashed rgba(102, 126, 234, 0.3);
    border-radius: 12px;
    transition: all 0.3s ease;
}
.image-upload:hover {
    border-color: rgba(102, 126, 234, 0.6);
    background: rgba(102, 126, 234, 0.05);
}
'''

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        # Header
        gr.Markdown("# 🎨 Qwen Image Edit - Katame Transfer")
        gr.Markdown("""
            Transform your images with AI-powered texture transfer using **Qwen Image Edit 2509**
            Powered by [2vXpSwA7/iroiro-lora](https://huggingface.co/2vXpSwA7/iroiro-lora) • [Qwen-Image-Lightning](https://huggingface.co/lightx2v/Qwen-Image-Lightning) ⚡
        """)

        gr.Markdown("---")

        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📥 Input Images")
                with gr.Row():
                    with gr.Column():
                        gr.Markdown("**🖼️ Content Image**")
                        content_image = gr.Image(label="", type="pil", elem_classes="image-upload")
                        with gr.Accordion("📁 Sample Images", open=False):
                            slotA_gallery = gr.Gallery(
                                value=slotA_images,
                                label="",
                                columns=3,
                                height="auto",
                                allow_preview=True,
                                show_label=False
                            )
                    with gr.Column():
                        gr.Markdown("**🎨 Texture Image**")
                        texture_image = gr.Image(label="", type="pil", elem_classes="image-upload")
                        with gr.Accordion("📁 Sample Textures", open=False):
                            slotB_gallery = gr.Gallery(
                                value=slotB_images,
                                label="",
                                columns=3,
                                height="auto",
                                allow_preview=True,
                                show_label=False
                            )

                gr.Markdown("### ✍️ Description")
                prompt = gr.Textbox(
                    label="",
                    info="",
                    placeholder="",
                    lines=2
                )

                button = gr.Button("✨ Generate Image", variant="primary", elem_classes="generate-btn")
                
                with gr.Accordion("⚙️ Advanced Settings", open=False):
                    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="🎯 Guidance Scale",
                        minimum=1.0,
                        maximum=10.0,
                        step=0.1,
                        value=1.0,
                        info="Higher values = stronger adherence to prompt"
                    )
                    num_inference_steps = gr.Slider(
                        label="⚡ Inference Steps",
                        minimum=1,
                        maximum=40,
                        step=1,
                        value=4,
                        info="More steps = higher quality (but slower)"
                    )

            with gr.Column(scale=1):
                gr.Markdown("### 🎭 Generated Result")
                output = gr.Image(label="", interactive=False, elem_classes="output-section")
                with gr.Row():
                    seed_display = gr.Number(label="🌱 Used Seed", interactive=False, visible=True)

        # Event handlers
        def select_slotA_image(evt: gr.SelectData):
            return slotA_images[evt.index]

        def select_slotB_image(evt: gr.SelectData):
            return slotB_images[evt.index]

        slotA_gallery.select(fn=select_slotA_image, outputs=content_image)
        slotB_gallery.select(fn=select_slotB_image, outputs=texture_image)

        button.click(
            fn=apply_texture,
            inputs=[
                content_image,
                texture_image,
                prompt,
                seed,
                randomize_seed,
                true_guidance_scale,
                num_inference_steps
            ],
            outputs=[output, seed_display]
        )

if __name__ == "__main__":
    demo.launch()