import gradio as gr
import numpy as np
import random
import torch
import spaces

from PIL import Image
from diffusers import FlowMatchEulerDiscreteScheduler, QwenImageEditPlusPipeline

import math

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

# Scheduler configuration for Lightning
scheduler_config = {
    "base_image_seq_len": 256,
    "base_shift": math.log(5),
    "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,
}

# Initialize scheduler with Lightning config
scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)

# Load the model pipeline
pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511", 
                                                 scheduler=scheduler,
                                                 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-fp32.safetensors"
)
pipe.fuse_lora()

# --- UI Constants and Helpers ---
MAX_SEED = np.iinfo(np.int32).max

def use_output_as_input(output_images):
    """Convert output images to input format for the gallery"""
    if output_images is None or len(output_images) == 0:
        return []
    return output_images

# --- Main Inference Function (with hardcoded negative prompt) ---
@spaces.GPU()
def infer(
    image_1,
    image_2,
    image_3,
    prompt,
    seed=42,
    randomize_seed=False,
    true_guidance_scale=1.0,
    num_inference_steps=4,
    height=None,
    width=None,
    num_images_per_prompt=1,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Run image-editing inference using the Qwen-Image-Edit pipeline.
    """
    
    # Hardcode the negative prompt as requested
    negative_prompt = " "
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    # Set up the generator for reproducibility
    generator = torch.Generator(device=device).manual_seed(seed)
    
    # Load input images into a list of PIL Images
    pil_images = []
    for item in [image_1, image_2, image_3]:
        if item is None: continue
        pil_images.append(item.convert("RGB"))

    if height==256 and width==256:
        height, width = None, None
    print(f"Calling pipeline with prompt: '{prompt}'")
    print(f"Negative Prompt: '{negative_prompt}'")
    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
    

    # Generate the image
    images = pipe(
        image=pil_images if len(pil_images) > 0 else None,
        prompt=prompt,
        height=height,
        width=width,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        generator=generator,
        true_cfg_scale=true_guidance_scale,
        num_images_per_prompt=num_images_per_prompt,
    ).images

    # Return images, seed, and make button visible
    return images[0], seed, gr.update(visible=True)


# ============================================
# 🎨 Comic Classic Theme - Toon Playground
# ============================================

css = """
/* ===== 🎨 Google Fonts Import ===== */
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@400;700&display=swap');

/* ===== 🎨 Comic Classic 배경 - 빈티지 페이퍼 + 도트 패턴 ===== */
.gradio-container {
    background-color: #FEF9C3 !important;
    background-image: 
        radial-gradient(#1F2937 1px, transparent 1px) !important;
    background-size: 20px 20px !important;
    min-height: 100vh !important;
    font-family: 'Comic Neue', cursive, sans-serif !important;
}

/* ===== 허깅페이스 상단 요소 숨김 ===== */
.huggingface-space-header,
#space-header,
.space-header,
[class*="space-header"],
.svelte-1ed2p3z,
.space-header-badge,
.header-badge,
[data-testid="space-header"],
.svelte-kqij2n,
.svelte-1ax1toq,
.embed-container > div:first-child {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    width: 0 !important;
    overflow: hidden !important;
    opacity: 0 !important;
    pointer-events: none !important;
}

/* ===== Footer 완전 숨김 ===== */
footer,
.footer,
.gradio-container footer,
.built-with,
[class*="footer"],
.gradio-footer,
.main-footer,
div[class*="footer"],
.show-api,
.built-with-gradio,
a[href*="gradio.app"],
a[href*="huggingface.co/spaces"] {
    display: none !important;
    visibility: hidden !important;
    height: 0 !important;
    padding: 0 !important;
    margin: 0 !important;
}

/* ===== 메인 컨테이너 ===== */
#col-container { 
    max-width: 1100px; 
    margin: 0 auto; 
}

/* ===== 🎨 헤더 타이틀 - 코믹 스타일 ===== */
.header-text h1 {
    font-family: 'Bangers', cursive !important;
    color: #1F2937 !important;
    font-size: 3.5rem !important;
    font-weight: 400 !important;
    text-align: center !important;
    margin-bottom: 0.5rem !important;
    text-shadow: 
        4px 4px 0px #FACC15,
        6px 6px 0px #1F2937 !important;
    letter-spacing: 3px !important;
    -webkit-text-stroke: 2px #1F2937 !important;
}

/* ===== 🎨 서브타이틀 ===== */
.subtitle {
    text-align: center !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.2rem !important;
    color: #1F2937 !important;
    margin-bottom: 1.5rem !important;
    font-weight: 700 !important;
}

/* ===== 🎨 카드/패널 - 만화 프레임 스타일 ===== */
.gr-panel,
.gr-box,
.gr-form,
.block,
.gr-group {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

/* ===== 🎨 입력 필드 (Textbox) ===== */
textarea, 
input[type="text"], 
input[type="number"] {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1rem !important;
    font-weight: 700 !important;
    transition: all 0.2s ease !important;
}

textarea:focus, 
input[type="text"]:focus, 
input[type="number"]:focus {
    border-color: #3B82F6 !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
    outline: none !important;
}

textarea::placeholder {
    color: #9CA3AF !important;
    font-weight: 400 !important;
}

/* ===== 🎨 Primary 버튼 - 코믹 블루 ===== */
.gr-button-primary,
button.primary,
.gr-button.primary,
.edit-btn {
    background: #3B82F6 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.3rem !important;
    letter-spacing: 2px !important;
    padding: 14px 28px !important;
    box-shadow: 5px 5px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-primary:hover,
button.primary:hover,
.gr-button.primary:hover,
.edit-btn:hover {
    background: #2563EB !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 7px 7px 0px #1F2937 !important;
}

.gr-button-primary:active,
button.primary:active,
.gr-button.primary:active,
.edit-btn:active {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary,
.use-output-btn {
    background: #EF4444 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    color: #FFFFFF !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.1rem !important;
    letter-spacing: 1px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
    transition: all 0.1s ease !important;
    text-shadow: 1px 1px 0px #1F2937 !important;
}

.gr-button-secondary:hover,
button.secondary:hover,
.use-output-btn:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

.gr-button-secondary:active,
button.secondary:active,
.use-output-btn:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 이미지 업로드/출력 영역 ===== */
.gr-image,
.image-container,
.image-upload {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    overflow: hidden !important;
    background: #FFFFFF !important;
}

/* ===== 🎨 아코디언 - 말풍선 스타일 ===== */
.gr-accordion {
    background: #FACC15 !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-accordion-header {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
}

/* ===== 🎨 슬라이더 스타일 ===== */
.gr-slider input[type="range"] {
    accent-color: #3B82F6 !important;
}

.gr-slider .gr-slider-label {
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    color: #1F2937 !important;
}

/* ===== 🎨 체크박스 스타일 ===== */
.gr-checkbox input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 20px !important;
    height: 20px !important;
}

.gr-checkbox label {
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    color: #1F2937 !important;
}

/* ===== 🎨 라벨 스타일 ===== */
label,
.gr-input-label,
.gr-block-label {
    color: #1F2937 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1rem !important;
}

span.gr-label {
    color: #1F2937 !important;
}

/* ===== 🎨 정보 텍스트 ===== */
.gr-info,
.info {
    color: #6B7280 !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 0.9rem !important;
}

/* ===== 🎨 프로그레스 바 ===== */
.progress-bar,
.gr-progress-bar {
    background: #3B82F6 !important;
    border: 2px solid #1F2937 !important;
    border-radius: 4px !important;
}

/* ===== 🎨 스크롤바 - 코믹 스타일 ===== */
::-webkit-scrollbar {
    width: 12px;
    height: 12px;
}

::-webkit-scrollbar-track {
    background: #FEF9C3;
    border: 2px solid #1F2937;
}

::-webkit-scrollbar-thumb {
    background: #3B82F6;
    border: 2px solid #1F2937;
    border-radius: 0px;
}

::-webkit-scrollbar-thumb:hover {
    background: #EF4444;
}

/* ===== 🎨 선택 하이라이트 ===== */
::selection {
    background: #FACC15;
    color: #1F2937;
}

/* ===== 🎨 링크 스타일 ===== */
a {
    color: #3B82F6 !important;
    text-decoration: none !important;
    font-weight: 700 !important;
}

a:hover {
    color: #EF4444 !important;
}

/* ===== 🎨 Row/Column 간격 ===== */
.gr-row {
    gap: 1.5rem !important;
}

.gr-column {
    gap: 1rem !important;
}

/* ===== 🎨 프롬프트 입력 영역 특별 스타일 ===== */
.prompt-input textarea {
    background: #EFF6FF !important;
    border: 4px dashed #3B82F6 !important;
    border-radius: 12px !important;
    font-size: 1.1rem !important;
    min-height: 120px !important;
}

.prompt-input textarea:focus {
    border-style: solid !important;
    background: #FFFFFF !important;
}

/* ===== 반응형 조정 ===== */
@media (max-width: 768px) {
    .header-text h1 {
        font-size: 2.2rem !important;
        text-shadow: 
            3px 3px 0px #FACC15,
            4px 4px 0px #1F2937 !important;
    }
    
    .gr-button-primary,
    button.primary {
        padding: 12px 20px !important;
        font-size: 1.1rem !important;
    }
    
    .gr-panel,
    .block {
        box-shadow: 4px 4px 0px #1F2937 !important;
    }
}

/* ===== 🎨 다크모드 비활성화 (코믹은 밝아야 함) ===== */
@media (prefers-color-scheme: dark) {
    .gradio-container {
        background-color: #FEF9C3 !important;
    }
}
"""

# --- Gradio UI Layout ---
with gr.Blocks(fill_height=True, css=css) as demo:
    gr.LoginButton(value="Option: HuggingFace 'Login' for extra GPU quota +", size="sm")       
    with gr.Column(elem_id="col-container"):
        
        # HOME Badge
        gr.HTML("""
            <div style="text-align: center; margin: 20px 0 10px 0;">
                <a href="https://www.humangen.ai" target="_blank" style="text-decoration: none;">
                    <img src="https://img.shields.io/static/v1?label=🏠 HOME&message=HUMANGEN.AI&color=0000ff&labelColor=ffcc00&style=for-the-badge" alt="HOME">
                </a>
            </div>
        """)
        
        # Header Title
        gr.Markdown(
            """
            # 🎨 QWEN IMAGE EDIT STUDIO 🖼️
            """,
            elem_classes="header-text"
        )
        
        gr.Markdown(
            """
            <p class="subtitle">✨ Upload an image and describe your edit - AI transforms it instantly! ✨</p>
            """,
        )

        with gr.Row(equal_height=False):
            # Left column - Input Images
            with gr.Column(scale=1, min_width=320):
                image_1 = gr.Image(
                    label="🖼️ Main Image",
                    type="pil",
                    interactive=True,
                    elem_classes="image-upload"
                )
                
                with gr.Accordion("📎 More Reference Images", open=False):
                    with gr.Row():
                        image_2 = gr.Image(
                            label="🖼️ Reference 2",
                            type="pil",
                            interactive=True
                        )
                        image_3 = gr.Image(
                            label="🖼️ Reference 3",
                            type="pil",
                            interactive=True
                        )

            # Right column - Output
            with gr.Column(scale=1, min_width=320):
                result = gr.Image(
                    label="🎯 Result",
                    type="pil",
                    interactive=False,
                    elem_classes="image-container"
                )
                
                use_output_btn = gr.Button(
                    "↗️ USE AS INPUT!",
                    variant="secondary",
                    size="sm",
                    visible=False,
                    elem_classes="use-output-btn"
                )
                
                gr.Markdown(
                    """
                    <p style="text-align: center; margin-top: 10px; font-weight: 700; color: #1F2937; font-family: 'Comic Neue', cursive;">
                        💡 Right-click on the image to save, or use the download button!
                    </p>
                    """
                )

        # Prompt Input Section
        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="✏️ Edit Instruction",
                    show_label=True,
                    placeholder="Describe the edit you want to make... e.g., 'Change the background to a beach sunset' or 'Add sunglasses to the person'",
                    lines=4,
                    elem_classes="prompt-input"
                )
                
                run_button = gr.Button(
                    "🚀 EDIT IMAGE!",
                    variant="primary",
                    size="lg",
                    elem_classes="edit-btn"
                )

        # Advanced Settings
        with gr.Accordion("⚙️ Advanced Settings", open=False):
            with gr.Row():
                seed = gr.Slider(
                    label="🎲 Seed",
                    minimum=0,
                    maximum=MAX_SEED,
                    step=1,
                    value=0,
                )
                randomize_seed = gr.Checkbox(
                    label="🔀 Randomize Seed",
                    value=True
                )

            with gr.Row():
                true_guidance_scale = gr.Slider(
                    label="🎯 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=20,
                )
            
            with gr.Row():
                height = gr.Slider(
                    label="📐 Height",
                    minimum=1024,
                    maximum=2048,
                    step=8,
                    value=1024,
                )
                width = gr.Slider(
                    label="📐 Width",
                    minimum=1024,
                    maximum=2048,
                    step=8,
                    value=1024,
                )

        # Tips Section
        gr.Markdown(
            """
            <div style="background: #FACC15; border: 3px solid #1F2937; border-radius: 8px; padding: 15px; margin-top: 20px; box-shadow: 4px 4px 0px #1F2937;">
                <h3 style="font-family: 'Bangers', cursive; color: #1F2937; margin: 0 0 10px 0; font-size: 1.3rem;">💡 PRO TIPS</h3>
                <ul style="font-family: 'Comic Neue', cursive; color: #1F2937; font-weight: 700; margin: 0; padding-left: 20px;">
                    <li>Be specific in your edit instructions for better results!</li>
                    <li>Use reference images for style transfer or composition guidance</li>
                    <li>Adjust inference steps: more steps = better quality, slower speed</li>
                    <li>Try different seeds if you don't like the result</li>
                </ul>
            </div>
            """
        )

    # Event Handlers
    gr.on(
        triggers=[run_button.click],
        fn=infer,
        inputs=[
            image_1,
            image_2,
            image_3,
            prompt,
            seed,
            randomize_seed,
            true_guidance_scale,
            num_inference_steps,
            height,
            width,
        ],
        outputs=[result, seed, use_output_btn],
    )

    use_output_btn.click(
        fn=use_output_as_input,
        inputs=[result],
        outputs=[image_1]
    )

if __name__ == "__main__":
    demo.launch(mcp_server=True, show_error=True)