import gradio as gr
from PIL import Image
import io
import zipfile

def resize_image_to_multiple_of_64_rgba(img: Image.Image):
    w, h = img.size
    # 计算最接近的 64 倍数（保底64）
    w64 = max(64, round(w / 64) * 64)
    h64 = max(64, round(h / 64) * 64)
    
    # 按最小缩放比，使原图能贴进 w64 x h64
    scale = min(w64 / w, h64 / h)
    new_width = int(w * scale)
    new_height = int(h * scale)
    
    # 用 RGBA、黑色背景(全不透明)
    background = Image.new("RGBA", (w64, h64), (0, 0, 0, 255))
    # 缩放原图到 new_width, new_height
    scaled_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 贴到背景中央，保留其透明度
    offset_x = (w64 - new_width) // 2
    offset_y = (h64 - new_height) // 2
    background.paste(scaled_img, (offset_x, offset_y), scaled_img)
    return background

def make_collage_2x2(four_rgba_images):
    # 4 张图，均已是 RGBA、同尺寸
    w, h = four_rgba_images[0].size
    # 新画布：2x2 => 宽度 2*w，高度 2*h，背景黑
    collage = Image.new("RGBA", (2 * w, 2 * h), (0, 0, 0, 255))
    collage.paste(four_rgba_images[0], (0, 0), four_rgba_images[0])
    collage.paste(four_rgba_images[1], (w, 0), four_rgba_images[1])
    collage.paste(four_rgba_images[2], (0, h), four_rgba_images[2])
    collage.paste(four_rgba_images[3], (w, h), four_rgba_images[3])
    
    cw, ch = collage.size
    # 若合并后任意边 > 2048，则等比例缩小
    if cw > 2048 or ch > 2048:
        scale = min(2048 / cw, 2048 / ch)
        new_cw = int(cw * scale)
        new_ch = int(ch * scale)
        collage = collage.resize((new_cw, new_ch), Image.Resampling.LANCZOS)
    
    return collage

def process_images_for_preview(uploaded_files):
    pil_images = []
    for f in uploaded_files:
        if f is not None:
            # 以 RGBA 读图，保证可保留透明通道
            img = Image.open(f.name).convert("RGBA")
            pil_images.append(img)

    results = []
    # 每 4 张为一组，不足 4 张跳过
    for i in range(0, len(pil_images), 4):
        group = pil_images[i : i + 4]
        if len(group) < 4:
            break
        # 每张做64倍数 resize + 黑色背景填充
        resized = [resize_image_to_multiple_of_64_rgba(im) for im in group]
        # 再 2×2 拼接
        collage = make_collage_2x2(resized)
        results.append(collage)
    
    return results

def process_and_zip_for_download(uploaded_files):
    collages = process_images_for_preview(uploaded_files)
    if not collages:
        # 没有任何拼接图，就返回 None，让界面不显示可下载链接
        return None
    
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
        for idx, img in enumerate(collages):
            # 转字节并写入 zip
            img_bytes = io.BytesIO()
            # 以 PNG 格式保存，带 RGBA
            img.save(img_bytes, format="PNG")
            img_bytes.seek(0)
            zf.writestr(f"collage_{idx+1}.png", img_bytes.read())
    buf.seek(0)
    return buf

with gr.Blocks() as demo:
    gr.Markdown("## 图片 2×2 拼接小工具")
    gr.Markdown(
        "1. 一次可上传多张图片，每 4 张为一组。\n"
        "2. 每组里 4 张图会先缩放到 64 的倍数、空余处用黑色填充，然后 2×2 拼接。\n"
        "3. 拼接结果若超过 2048×2048，则缩小到不超过 2048。\n"
        "4. 不足 4 张的剩余图片忽略。\n"
        "5. 支持保留 PNG 透明度，拼图时使用黑色背景。"
    )
    
    with gr.Row():
        with gr.Column():
            file_input = gr.Files(label="上传图片（可多选）", file_types=["image"])
            preview_btn = gr.Button("生成预览")
            download_btn = gr.Button("打包下载 ZIP")
        with gr.Column():
            # Gallery 构造上不使用 .style() 以兼容老版本
            gallery_output = gr.Gallery(label="拼接结果预览", columns=2)
            zip_output = gr.File(label="下载拼接结果 ZIP", interactive=False, visible=False)
    
    # 点击“生成预览” -> 返回拼接图列表
    preview_btn.click(
        fn=process_images_for_preview,
        inputs=[file_input],
        outputs=[gallery_output]
    )
    
    # 点击“打包下载 ZIP” -> 生成 zip
    download_btn.click(
        fn=process_and_zip_for_download,
        inputs=[file_input],
        outputs=[zip_output]
    )

demo.launch()