import os
import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers.pipelines.glm_image import GlmImagePipeline
from PIL import Image

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

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

pipe = GlmImagePipeline.from_pretrained(
    "zai-org/GLM-Image",
    torch_dtype=torch.bfloat16,
).to("cuda")


@spaces.GPU(duration=120)
def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, 
          num_inference_steps=50, guidance_scale=1.5, progress=gr.Progress(track_tqdm=True)):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    width = (width // 32) * 32
    height = (height // 32) * 32
    
    generator = torch.Generator(device="cuda").manual_seed(seed)

    image_list = None
    if input_images is not None and len(input_images) > 0:
        image_list = []
        for item in input_images:
            img = item[0] if isinstance(item, tuple) else item
            if isinstance(img, str):
                img = Image.open(img).convert("RGB")
            elif isinstance(img, Image.Image):
                img = img.convert("RGB")
            image_list.append(img)

    pipe_kwargs = {
        "prompt": prompt,
        "height": height,
        "width": width,
        "num_inference_steps": num_inference_steps,
        "guidance_scale": guidance_scale,
        "generator": generator,
    }

    if image_list is not None:
        pipe_kwargs["image"] = image_list

    image = pipe(**pipe_kwargs).images[0]
    
    return image, seed


def update_dimensions_from_image(image_list):
    if image_list is None or len(image_list) == 0:
        return 1024, 1024
    
    item = image_list[0]
    img = item[0] if isinstance(item, tuple) else item
    
    if isinstance(img, str):
        img = Image.open(img)
    
    img_width, img_height = img.size
    aspect_ratio = img_width / img_height
    
    if aspect_ratio >= 1:
        new_width = 1024
        new_height = int(1024 / aspect_ratio)
    else:
        new_height = 1024
        new_width = int(1024 * aspect_ratio)
    
    new_width = round(new_width / 32) * 32
    new_height = round(new_height / 32) * 32
    
    new_width = max(256, min(MAX_IMAGE_SIZE, new_width))
    new_height = max(256, min(MAX_IMAGE_SIZE, new_height))
    
    return new_width, new_height

css = """
/* POP ART Style */
@import url('https://fonts.googleapis.com/css2?family=Bangers&family=Comic+Neue:wght@700&display=swap');

* {
    font-family: 'Comic Neue', cursive !important;
}

body, .gradio-container {
    background: #ffeb3b !important;
    min-height: 100vh;
}

.gradio-container {
    background: 
        radial-gradient(circle at 10px 10px, #ff1744 4px, transparent 4px),
        radial-gradient(circle at 30px 30px, #ff1744 3px, transparent 3px),
        linear-gradient(135deg, #ffeb3b 0%, #fff176 100%) !important;
    background-size: 40px 40px, 40px 40px, 100% 100% !important;
}

#col-container {
    margin: 0 auto;
    max-width: 1150px;
    padding: 25px 20px;
}

/* Header - Comic Speech Bubble */
.header-box {
    background: #ffffff;
    border: 5px solid #000000;
    border-radius: 30px;
    padding: 25px 30px;
    text-align: center;
    margin-bottom: 25px;
    box-shadow: 8px 8px 0 #000000;
    position: relative;
}

.header-box::after {
    content: "";
    position: absolute;
    bottom: -25px;
    left: 50px;
    width: 0;
    height: 0;
    border: 20px solid transparent;
    border-top-color: #000000;
    border-bottom: 0;
}

.header-box::before {
    content: "";
    position: absolute;
    bottom: -17px;
    left: 53px;
    width: 0;
    height: 0;
    border: 16px solid transparent;
    border-top-color: #ffffff;
    border-bottom: 0;
    z-index: 1;
}

.badge-row {
    display: flex;
    justify-content: center;
    margin-bottom: 12px;
}

.title-text {
    font-family: 'Bangers', cursive !important;
    font-size: 3rem !important;
    color: #ff1744 !important;
    margin: 8px 0 !important;
    text-shadow: 
        3px 3px 0 #000000,
        -1px -1px 0 #000000,
        1px -1px 0 #000000,
        -1px 1px 0 #000000;
    letter-spacing: 3px;
}

.subtitle-text {
    font-family: 'Comic Neue', cursive !important;
    color: #1565c0 !important;
    font-size: 1.1rem !important;
    font-weight: 700 !important;
    margin: 0 !important;
    text-transform: uppercase;
}

/* Cards - Bold Comic Style */
.pop-card {
    background: #ffffff !important;
    border: 4px solid #000000 !important;
    border-radius: 20px !important;
    padding: 20px !important;
    box-shadow: 6px 6px 0 #000000 !important;
    position: relative;
}

.gr-panel, .gr-box, .gr-form, .block, .form {
    background: #ffffff !important;
    border: 3px solid #000000 !important;
    border-radius: 15px !important;
    box-shadow: 4px 4px 0 #000000 !important;
}

/* Prompt Input - Speech Bubble */
.prompt-box textarea {
    background: #ffffff !important;
    border: 4px solid #000000 !important;
    border-radius: 20px !important;
    color: #000000 !important;
    font-size: 1.1rem !important;
    font-weight: 700 !important;
    padding: 18px !important;
    min-height: 90px !important;
    transition: all 0.2s ease !important;
    box-shadow: 4px 4px 0 #000000 !important;
}

.prompt-box textarea:focus {
    border-color: #ff1744 !important;
    box-shadow: 6px 6px 0 #ff1744 !important;
    transform: translate(-2px, -2px);
}

.prompt-box textarea::placeholder {
    color: #757575 !important;
    font-weight: 600 !important;
}

/* Generate Button - POW! Style */
.pow-btn {
    background: linear-gradient(180deg, #ff1744 0%, #d50000 100%) !important;
    border: 4px solid #000000 !important;
    border-radius: 15px !important;
    color: #ffffff !important;
    font-family: 'Bangers', cursive !important;
    font-weight: 400 !important;
    font-size: 1.5rem !important;
    letter-spacing: 2px;
    padding: 16px 35px !important;
    transition: all 0.15s ease !important;
    box-shadow: 5px 5px 0 #000000 !important;
    text-shadow: 2px 2px 0 #000000;
    text-transform: uppercase;
}

.pow-btn:hover {
    transform: translate(-3px, -3px) !important;
    box-shadow: 8px 8px 0 #000000 !important;
    background: linear-gradient(180deg, #ff5252 0%, #ff1744 100%) !important;
}

.pow-btn:active {
    transform: translate(2px, 2px) !important;
    box-shadow: 2px 2px 0 #000000 !important;
}

/* Result Image - Comic Panel */
.result-box {
    background: #ffffff !important;
    border: 5px solid #000000 !important;
    border-radius: 20px !important;
    padding: 15px !important;
    box-shadow: 8px 8px 0 #000000 !important;
    min-height: 450px !important;
    position: relative;
}

.result-box::before {
    content: "★ OUTPUT ★";
    position: absolute;
    top: -15px;
    left: 20px;
    background: #ffeb3b;
    border: 3px solid #000000;
    padding: 3px 15px;
    font-family: 'Bangers', cursive !important;
    font-size: 1rem;
    color: #000000;
    z-index: 10;
}

.result-box img {
    border-radius: 12px !important;
    border: 3px solid #000000 !important;
}

/* Accordion - Comic Panels */
.gr-accordion {
    background: #e3f2fd !important;
    border: 3px solid #000000 !important;
    border-radius: 15px !important;
    margin-top: 20px !important;
    overflow: hidden !important;
    box-shadow: 4px 4px 0 #000000 !important;
}

.gr-accordion-header {
    background: #2196f3 !important;
    color: #ffffff !important;
    font-weight: 700 !important;
    padding: 12px 18px !important;
    border-bottom: 3px solid #000000 !important;
    text-transform: uppercase;
}

/* Gallery */
.gallery-box {
    background: #fff9c4 !important;
    border-radius: 12px !important;
    padding: 12px !important;
    border: 3px dashed #000000 !important;
}

.gallery-box img {
    border-radius: 10px !important;
    border: 3px solid #000000 !important;
}

/* Labels */
label, .label-wrap, span {
    color: #000000 !important;
    font-weight: 700 !important;
    text-transform: uppercase;
    font-size: 0.9rem !important;
}

/* Sliders */
input[type="range"] {
    accent-color: #ff1744 !important;
    height: 8px !important;
}

input[type="range"]::-webkit-slider-track {
    background: #000000 !important;
    border-radius: 5px !important;
    height: 8px !important;
}

input[type="range"]::-webkit-slider-thumb {
    background: #ffeb3b !important;
    border: 3px solid #000000 !important;
    width: 20px !important;
    height: 20px !important;
}

/* Number Input */
input[type="number"] {
    background: #ffffff !important;
    border: 3px solid #000000 !important;
    border-radius: 10px !important;
    color: #000000 !important;
    font-weight: 700 !important;
    box-shadow: 3px 3px 0 #000000 !important;
}

/* Checkbox */
input[type="checkbox"] {
    accent-color: #ff1744 !important;
    width: 20px !important;
    height: 20px !important;
    border: 3px solid #000000 !important;
}

/* Scrollbar */
::-webkit-scrollbar {
    width: 12px;
}

::-webkit-scrollbar-track {
    background: #ffeb3b;
    border: 2px solid #000000;
}

::-webkit-scrollbar-thumb {
    background: #ff1744;
    border: 2px solid #000000;
    border-radius: 0;
}

/* Layout */
.main-row {
    gap: 30px !important;
    align-items: flex-start !important;
}

.input-col {
    flex: 1 !important;
}

.output-col {
    flex: 1.1 !important;
}

/* Halftone Pattern Overlay */
.halftone-overlay {
    position: fixed;
    top: 0;
    left: 0;
    width: 100%;
    height: 100%;
    background-image: radial-gradient(#000000 1px, transparent 1px);
    background-size: 4px 4px;
    opacity: 0.03;
    pointer-events: none;
    z-index: 9999;
}

/* Action Words */
.action-word {
    font-family: 'Bangers', cursive !important;
    color: #ff1744;
    text-shadow: 2px 2px 0 #000000;
}
"""

with gr.Blocks(css=css, theme=gr.themes.Base()) as demo:
    
    # Halftone overlay
    gr.HTML('<div class="halftone-overlay"></div>')
    
    with gr.Column(elem_id="col-container"):
        # Header - Comic Speech Bubble
        gr.HTML("""
            <div class="header-box">
                <div class="badge-row">
                    <a href="https://www.humangen.ai" target="_blank">
                        <img src="https://img.shields.io/static/v1?label=Free%20AI%20HUB&message=humangen.ai&color=%23ff1744&labelColor=%23000000&logo=huggingface&logoColor=white&style=for-the-badge" alt="badge">
                    </a>
                </div>
                <h1 class="title-text">💥 GLM-IMAGE GENERATOR</h1>
                <p class="subtitle-text">⚡ 9B hybrid AI model for text-accurate image generation ⚡</p>
            </div>
        """)
        
        with gr.Row(elem_classes="main-row"):
            # Left - Input
            with gr.Column(elem_classes="input-col"):
                with gr.Group(elem_classes="pop-card"):
                    prompt = gr.Text(
                        label="Prompt",
                        show_label=False,
                        max_lines=4,
                        placeholder="💬 Type your image idea here...",
                        container=False,
                        elem_classes="prompt-box"
                    )
                    
                    run_button = gr.Button("⚡ POW! GENERATE ⚡", variant="primary", elem_classes="pow-btn")
                
                with gr.Accordion("📸 INPUT IMAGES", open=False):
                    input_images = gr.Gallery(
                        label="Upload images",
                        type="pil",
                        columns=3,
                        rows=1,
                        elem_classes="gallery-box"
                    )
                
                with gr.Accordion("🔧 SETTINGS", open=False):
                    seed = gr.Slider(label="SEED", minimum=0, maximum=MAX_SEED, step=1, value=42)
                    randomize_seed = gr.Checkbox(label="RANDOMIZE", value=True)
                    
                    with gr.Row():
                        width = gr.Slider(label="WIDTH", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                        height = gr.Slider(label="HEIGHT", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
                    
                    with gr.Row():
                        num_inference_steps = gr.Slider(label="STEPS", minimum=1, maximum=100, step=1, value=50)
                        guidance_scale = gr.Slider(label="GUIDANCE", minimum=0.0, maximum=10.0, step=0.1, value=1.5)
            
            # Right - Output
            with gr.Column(elem_classes="output-col"):
                result = gr.Image(label="Result", show_label=False, elem_classes="result-box")

    input_images.upload(
        fn=update_dimensions_from_image,
        inputs=[input_images],
        outputs=[width, height]
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale],
        outputs=[result, seed]
    )

demo.launch()