import gradio as gr
from PIL import Image
import io
import zipfile
import random
import tempfile
import os

def random_black_or_white():
    return (0, 0, 0, 255) if random.random() < 0.5 else (255, 255, 255, 255)

def random_non_black_white():
    while True:
        r = random.randint(0, 255)
        g = random.randint(0, 255)
        b = random.randint(0, 255)
        # 不要纯黑纯白
        if not (r == g == b == 0 or r == g == b == 255):
            return (r, g, b, 255)

def limit_2048(img: Image.Image):
    w, h = img.size
    if w > 2048 or h > 2048:
        scale = min(2048 / w, 2048 / h)
        nw = int(w * scale)
        nh = int(h * scale)
        img = img.resize((nw, nh), Image.Resampling.LANCZOS)
    return img

def resize_to_64_multiple(img: Image.Image):
    w, h = img.size
    w64 = max(64, round(w / 64) * 64)
    h64 = max(64, round(h / 64) * 64)
    
    scale = min(w64 / w, h64 / h)
    nw = int(w * scale)
    nh = int(h * scale)
    
    bg_color = random_black_or_white()
    background = Image.new("RGBA", (w64, h64), bg_color)
    scaled = img.resize((nw, nh), Image.Resampling.LANCZOS)
    ox = (w64 - nw) // 2
    oy = (h64 - nh) // 2
    background.paste(scaled, (ox, oy), scaled)
    return background

def make_collage_2x2(images_4):
    w, h = images_4[0].size
    collage = Image.new("RGBA", (2*w, 2*h), (0,0,0,255))
    collage.paste(images_4[0], (0,0), images_4[0])
    collage.paste(images_4[1], (w,0), images_4[1])
    collage.paste(images_4[2], (0,h), images_4[2])
    collage.paste(images_4[3], (w,h), images_4[3])
    return limit_2048(collage)

def make_collage_leftover(images_leftover):
    n = len(images_leftover)
    if n < 1 or n > 3:
        return None
    
    resized_list = [resize_to_64_multiple(img) for img in images_leftover]
    max_w = max(im.size[0] for im in resized_list)
    max_h = max(im.size[1] for im in resized_list)
    
    uniformed = []
    for rimg in resized_list:
        w, h = rimg.size
        if (w,h) == (max_w, max_h):
            uniformed.append(rimg)
        else:
            bg_color = rimg.getpixel((0,0))
            bg = Image.new("RGBA", (max_w, max_h), bg_color)
            offx = (max_w - w)//2
            offy = (max_h - h)//2
            bg.paste(rimg, (offx, offy), rimg)
            uniformed.append(bg)
    
    if n == 1:
        possible_layouts = [(1,1), (1,2), (2,1), (2,2)]
    elif n == 2:
        possible_layouts = [(1,2), (2,1), (2,2)]
    else: # n == 3
        possible_layouts = [(2,2)]
    
    rows, cols = random.choice(possible_layouts)
    big_w = cols * max_w
    big_h = rows * max_h
    collage = Image.new("RGBA", (big_w, big_h), (0,0,0,255))
    
    cells = [(r, c) for r in range(rows) for c in range(cols)]
    random.shuffle(cells)
    
    for i, img_ in enumerate(uniformed):
        r, c = cells[i]
        ox = c * max_w
        oy = r * max_h
        collage.paste(img_, (ox, oy), img_)
    
    leftover_cells = cells[n:]
    for (r, c) in leftover_cells:
        color_ = random_non_black_white()
        rect = Image.new("RGBA", (max_w, max_h), color_)
        collage.paste(rect, (c*max_w, r*max_h), rect)
    
    return limit_2048(collage)

def process_images(uploaded_files):
    pil_images = []
    for f in uploaded_files:
        if f is not None:
            img = Image.open(f.name).convert("RGBA")
            pil_images.append(img)
    
    results = []
    total = len(pil_images)
    groups_4 = total // 4
    leftover = total % 4
    
    idx = 0
    for _ in range(groups_4):
        group_4 = pil_images[idx:idx+4]
        idx += 4
        resized_4 = [resize_to_64_multiple(im) for im in group_4]
        max_w = max(im.size[0] for im in resized_4)
        max_h = max(im.size[1] for im in resized_4)
        
        final_4 = []
        for rimg in resized_4:
            w, h = rimg.size
            if (w,h) == (max_w, max_h):
                final_4.append(rimg)
            else:
                bg_color = rimg.getpixel((0,0))
                bg = Image.new("RGBA", (max_w, max_h), bg_color)
                offx = (max_w - w)//2
                offy = (max_h - h)//2
                bg.paste(rimg, (offx, offy), rimg)
                final_4.append(bg)
        
        collage_2x2 = make_collage_2x2(final_4)
        results.append(collage_2x2)
    
    if leftover > 0:
        leftover_imgs = pil_images[idx:]
        collage_left = make_collage_leftover(leftover_imgs)
        if collage_left is not None:
            results.append(collage_left)
    
    return results

def make_zip(uploaded_files):
    """
    真正构建ZIP，并返回临时文件路径（不是BytesIO），
    以防 Gradio 把 BytesIO 当成字符串解析而报错。
    """
    collages = process_images(uploaded_files)
    if not collages:
        return None  # 让前端做判断
    
    # 在内存里打包
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        for i, img in enumerate(collages, start=1):
            img_bytes = io.BytesIO()
            img.save(img_bytes, format="PNG")
            img_bytes.seek(0)
            zf.writestr(f"collage_{i}.png", img_bytes.read())
    zip_buffer.seek(0)
    
    # 将 BytesIO 写入临时文件，再返回临时文件路径
    with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
        tmp.write(zip_buffer.getvalue())
        tmp_path = tmp.name
    
    return tmp_path

def on_zip_click(files):
    """
    返回 (zip_file_path, message) 给两个 output:
    1. gr.File 要求返回一个str路径或None
    2. gr.Textbox 用来输出提示信息
    """
    path = make_zip(files)
    if path is None:
        return (None, "无可下载内容 - 可能没上传图片或无法拼接")
    else:
        return (path, "打包完成！点击上方链接下载ZIP")

with gr.Blocks() as demo:
    gr.Markdown("## 2×2 拼接小工具（兼容不足4张、随机填充、保留透明）")
    with gr.Row():
        with gr.Column():
            file_input = gr.Files(label="上传多张图片(可多选)", file_types=["image"])
            preview_btn = gr.Button("生成预览")
            zip_btn = gr.Button("打包下载 ZIP")
        with gr.Column():
            gallery_out = gr.Gallery(label="拼接结果预览", columns=2)
            # 注意：必须 visible=True，这样点击后可以直接更新
            zip_file_out = gr.File(label="下载拼接结果 ZIP", visible=True, interactive=False)
            msg_out = gr.Textbox(label="提示信息", interactive=False)
    
    preview_btn.click(
        fn=process_images,
        inputs=[file_input],
        outputs=[gallery_out]
    )
    zip_btn.click(
        fn=on_zip_click,
        inputs=[file_input],
        outputs=[zip_file_out, msg_out]
    )

demo.launch()