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

def resize_image_to_multiple_of_64(img: Image.Image, fill_color=(0, 0, 0)):
    """
    将单张图片缩放到 (w', h')，其中 w'、h' 均为 64 的整数倍。
    若原图宽高比与 w'/h' 不完全一致，则在空余处用 fill_color 填充（不裁切）。
    """
    w, h = img.size
    
    # 计算最接近的 64 倍数
    # 若 round(...) 结果为 0，至少保底 64
    w64 = max(64, round(w / 64) * 64)
    h64 = max(64, round(h / 64) * 64)
    
    # 根据最小缩放策略，保证图片能放进 w64 × h64
    scale = min(w64 / w, h64 / h)
    new_width = int(w * scale)
    new_height = int(h * scale)
    
    # 创建背景
    new_img = Image.new("RGB", (w64, h64), fill_color)
    # 将原图缩放
    scaled_img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    # 居中贴到背景上
    offset_x = (w64 - new_width) // 2
    offset_y = (h64 - new_height) // 2
    new_img.paste(scaled_img, (offset_x, offset_y))
    return new_img

def make_collage(four_images, fill_color=(0, 0, 0)):
    """
    将 4 张同样尺寸的图片（2×2）拼接成一张大图。
    如果结果大于 2048×2048，则等比例缩小到不超过 2048。
    """
    assert len(four_images) == 4, "必须是 4 张图"
    w, h = four_images[0].size  # 假设 4 张图尺寸相同
    collage = Image.new("RGB", (2 * w, 2 * h), fill_color)
    # 依次贴上
    collage.paste(four_images[0], (0, 0))
    collage.paste(four_images[1], (w, 0))
    collage.paste(four_images[2], (0, h))
    collage.paste(four_images[3], (w, h))
    
    cw, ch = collage.size
    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(uploaded_files):
    """
    处理用户上传的多张图片：
    1. 每组 4 张图 -> 统一到 64 的倍数尺寸 -> 拼成 2×2 的大图
    2. 大于 2048×2048 则缩小
    3. 返回多张拼接图（以列表的形式）供 Gradio 显示
    4. 若最后一组不足 4 张，则跳过
    """
    # 读取成 PIL 对象
    pil_images = []
    for file_obj in uploaded_files:
        if file_obj is not None:
            img = Image.open(file_obj.name).convert("RGB")
            pil_images.append(img)
    
    results = []
    # 以 4 张为一组进行处理
    for i in range(0, len(pil_images), 4):
        group = pil_images[i : i + 4]
        if len(group) < 4:
            # 不足 4 张，跳过
            break
        
        # 分别 resize
        resized = [resize_image_to_multiple_of_64(img, fill_color=(0,0,0)) for img in group]
        # 拼接
        collage = make_collage(resized, fill_color=(0,0,0))
        results.append(collage)
    
    return results

def process_and_zip(uploaded_files):
    """
    在内存中创建 Zip，将所有生成的拼接图打包。
    返回一个 BytesIO 给 Gradio 的 File 输出，用于下载。
    """
    # 先按需求生成拼接结果
    collages = process_images(uploaded_files)
    
    # 如果没有生成任何拼接图，则返回 None，让前端不显示下载
    if not collages:
        return None
    
    # 在内存中创建 zip
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as zf:
        for idx, img in enumerate(collages):
            # 将 PIL 图像转成字节并写入 zip
            img_bytes = io.BytesIO()
            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

########################################
#        构建 Gradio 界面
########################################
css_custom = """
/* 你可根据需要，在此自定义或覆盖 Gradio 的 CSS 样式 */
"""

with gr.Blocks(css=css_custom) as demo:
    gr.Markdown("## 图片 2×2 拼接小工具")
    gr.Markdown(
        "1. 请选择若干张图片（数量需为 4 的倍数才能完整拼接）。<br>"
        "2. 每 4 张会被缩放并以 2×2 的形式合并成新图。<br>"
        "3. 如果合并后图像尺寸超过 2048×2048，会自动缩小。<br>"
        "4. 你可以在下方预览结果或打包下载。"
    )
    
    with gr.Row():
        with gr.Column():
            file_input = gr.Files(label="上传图片（可多选）", file_types=["image"])
            btn_preview = gr.Button("生成预览")
            btn_zip = gr.Button("打包下载 ZIP")
        
        with gr.Column():
            gallery_output = gr.Gallery(label="拼接结果预览").style(grid=[2], height="auto")
            zip_file = gr.File(label="下载拼接结果 ZIP", interactive=False, visible=False)
    
    # 点击“生成预览” -> 显示拼接好的图片（可能多张）
    btn_preview.click(
        fn=process_images,
        inputs=[file_input],
        outputs=[gallery_output]
    )
    
    # 点击“打包下载 ZIP” -> 生成 zip 文件供下载
    btn_zip.click(
        fn=process_and_zip,
        inputs=[file_input],
        outputs=[zip_file]
    )

demo.launch()