import os
import uuid
import numpy as np
import random
import tempfile
import spaces
import zipfile 
from PIL import Image
from diffusers import QwenImageLayeredPipeline
import torch
from pptx import Presentation
import gradio as gr


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

from huggingface_hub import login
login(token=os.environ.get('hf'))

dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = QwenImageLayeredPipeline.from_pretrained("Qwen/Qwen-Image-Layered", torch_dtype=dtype).to(device)

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 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

def export_gallery(images):
    images = [e[0] for e in images]
    pptx_path = imagelist_to_pptx(images)
    return pptx_path

def export_gallery_zip(images):
    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

@spaces.GPU(duration=180)
def infer(input_image,
          seed=777,
          randomize_seed=False,
          prompt=None,
          neg_prompt=" ",
          true_guidance_scale=4.0,
          num_inference_steps=50,
          layer=4,
          cfg_norm=True,
          use_en_prompt=True):
    
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
        
    if isinstance(input_image, list):
        input_image = input_image[0]
        
    if isinstance(input_image, str):
        pil_image = Image.open(input_image).convert("RGB").convert("RGBA")
    elif isinstance(input_image, Image.Image):
        pil_image = input_image.convert("RGB").convert("RGBA")
    elif isinstance(input_image, np.ndarray):
        pil_image = Image.fromarray(input_image).convert("RGB").convert("RGBA")
    else:
        raise ValueError("Unsupported input_image type: %s" % type(input_image))
    
    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": num_inference_steps,
        "num_images_per_prompt": 1,
        "layers": layer,
        "resolution": 640,
        "cfg_normalize": cfg_norm,
        "use_en_prompt": use_en_prompt, 
    }
    print(inputs)
    with torch.inference_mode():
        output = pipeline(**inputs)
        output_images = output.images[0]
    
    output = []
    temp_files = []
    for i, image in enumerate(output_images):
        output.append(image)
        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        image.save(tmp.name)
        temp_files.append(tmp.name)
    
    pptx_path = imagelist_to_pptx(temp_files)
    
    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(temp_files):
                zipf.write(img_path, f"layer_{i+1}.png")
        zip_path = tmp.name
    
    # 정보 로그 생성
    info_log = f"""✅ DECOMPOSITION COMPLETE!
{'=' * 50}
🖼️ Input Image Info:
   • Size: {pil_image.size[0]} x {pil_image.size[1]}
   • Mode: {pil_image.mode}
{'=' * 50}
⚙️ Generation Settings:
   • Seed: {seed}
   • Layers: {layer}
   • Steps: {num_inference_steps}
   • CFG Scale: {true_guidance_scale}
{'=' * 50}
📦 Output:
   • Generated Layers: {len(output_images)}
   • PPTX: Ready to download!
   • ZIP: Ready to download!
{'=' * 50}
💾 All files ready for download!"""
    
    return output, pptx_path, zip_path, info_log

ensure_dirname(LOG_DIR)
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",
]


# ============================================
# 🎨 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: 1200px; 
    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;
    transition: all 0.2s ease !important;
}

.gr-panel:hover,
.block:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 8px 8px 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 {
    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 {
    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 {
    transform: translate(3px, 3px) !important;
    box-shadow: 2px 2px 0px #1F2937 !important;
}

/* ===== 🎨 Secondary 버튼 - 코믹 레드 ===== */
.gr-button-secondary,
button.secondary,
.decompose-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,
.decompose-btn:hover {
    background: #DC2626 !important;
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

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

/* ===== 🎨 로그 출력 영역 ===== */
.info-log textarea {
    background: #1F2937 !important;
    color: #10B981 !important;
    font-family: 'Courier New', monospace !important;
    font-size: 0.9rem !important;
    font-weight: 400 !important;
    border: 3px solid #10B981 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #10B981 !important;
}

/* ===== 🎨 이미지 업로드 영역 ===== */
.image-upload {
    border: 4px dashed #3B82F6 !important;
    border-radius: 12px !important;
    background: #EFF6FF !important;
    transition: all 0.2s ease !important;
}

.image-upload:hover {
    border-color: #EF4444 !important;
    background: #FEF2F2 !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-gallery,
.gallery-container {
    border: 4px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 8px 8px 0px #1F2937 !important;
    overflow: hidden !important;
    background: #FFFFFF !important;
}

.gr-gallery .thumbnail-item {
    border: 3px solid #1F2937 !important;
    border-radius: 6px !important;
    transition: all 0.2s ease !important;
}

.gr-gallery .thumbnail-item:hover {
    transform: scale(1.05) !important;
    box-shadow: 4px 4px 0px #3B82F6 !important;
}

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

/* ===== 🎨 파일 다운로드 영역 ===== */
.gr-file,
.file-container {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 4px 4px 0px #1F2937 !important;
}

.gr-file:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
}

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

.gr-slider {
    background: #FFFFFF !important;
}

/* ===== 🎨 체크박스 스타일 ===== */
input[type="checkbox"] {
    accent-color: #3B82F6 !important;
    width: 20px !important;
    height: 20px !important;
    border: 2px solid #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;
}

/* ===== 🎨 Examples 섹션 ===== */
.gr-examples {
    background: #FFFFFF !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    box-shadow: 6px 6px 0px #1F2937 !important;
    padding: 1rem !important;
}

.gr-examples .gr-sample {
    border: 2px solid #1F2937 !important;
    border-radius: 6px !important;
    transition: all 0.2s ease !important;
}

.gr-examples .gr-sample:hover {
    transform: translate(-2px, -2px) !important;
    box-shadow: 4px 4px 0px #3B82F6 !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;
}

/* ===== 🎨 다운로드 버튼 강조 ===== */
.download-section {
    background: linear-gradient(135deg, #FACC15 0%, #FEF9C3 100%) !important;
    border: 3px solid #1F2937 !important;
    border-radius: 8px !important;
    padding: 1rem !important;
    box-shadow: 4px 4px 0px #1F2937 !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;
    }
}
"""

# Build the Gradio interface
with gr.Blocks(fill_height=True, css=css) as demo:
    
    # 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 LAYERED DECOMPOSER 🖼️
        """,
        elem_classes="header-text"
    )
    
    gr.Markdown(
        """
        <p class="subtitle">🔮 Upload an image and decompose it into magical layers! ✨ PPTX & ZIP export ready! 📦</p>
        """,
    )
    
    with gr.Row(equal_height=False):
        # Left column - Input
        with gr.Column(scale=1, min_width=350):
            input_image = gr.Image(
                label="🖼️ Upload Your Image",
                image_mode="RGBA",
                elem_classes="image-upload"
            )
            
            run_button = gr.Button(
                "🎨 DECOMPOSE INTO LAYERS! 🔮",
                variant="primary",
                size="lg",
                elem_classes="decompose-btn"
            )
            
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                prompt = gr.Textbox(
                    label="✏️ Prompt (Optional)",
                    placeholder="Describe the image content including hidden elements...",
                    value="",
                    lines=2,
                )
                neg_prompt = gr.Textbox(
                    label="🚫 Negative Prompt (Optional)",
                    placeholder="What to avoid in generation...",
                    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="🔄 Inference Steps",
                    minimum=1,
                    maximum=50,
                    step=1,
                    value=50,
                )

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

                cfg_norm = gr.Checkbox(label="✅ Enable CFG Normalization", value=True)
                use_en_prompt = gr.Checkbox(label="🌐 Auto Caption (EN=True, ZH=False)", value=True)
            
            with gr.Accordion("📜 Processing Log", open=True):
                info_log = gr.Textbox(
                    label="",
                    placeholder="Upload an image and click decompose to see info...",
                    lines=14,
                    max_lines=20,
                    interactive=False,
                    elem_classes="info-log"
                )
        
        # Right column - Output
        with gr.Column(scale=2, min_width=500):
            gallery = gr.Gallery(
                label="🎭 Decomposed Layers",
                columns=4,
                rows=2,
                format="png",
                height=400
            )
            
            gr.Markdown(
                """
                <p style="text-align: center; margin: 15px 0; font-weight: 700; color: #1F2937; font-size: 1.1rem;">
                    💾 Download Your Layers Below! 👇
                </p>
                """
            )
            
            with gr.Row(elem_classes="download-section"):
                export_file = gr.File(label="📊 Download PPTX")
                export_zip_file = gr.File(label="📦 Download ZIP")
            
            gr.Markdown(
                """
                <p style="text-align: center; margin-top: 15px; font-weight: 700; color: #6B7280;">
                    💡 PPTX preserves layers for editing • ZIP contains all PNG files
                </p>
                """
            )

    # Examples Section
    gr.Markdown(
        """
        <p style="text-align: center; margin: 20px 0 10px 0; font-family: 'Bangers', cursive; font-size: 1.5rem; color: #1F2937;">
            🌟 TRY THESE EXAMPLES! 🌟
        </p>
        """
    )
    
    gr.Examples(
        examples=examples,
        inputs=[input_image], 
        outputs=[gallery, export_file, export_zip_file, info_log],
        fn=infer, 
        examples_per_page=14,
        cache_examples=False,
        run_on_click=True
    )

    # Connect the button
    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, info_log],
    )

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