import io
from PIL import Image


def merge_images(images, direction="horizontal", gap=0, bg_color=(255, 255, 255)):
    """拼接多张图片"""
    if not images:
        raise ValueError("没有图片")

    imgs = [Image.open(f) for f in images]

    if direction == "horizontal":
        total_w = sum(im.width for im in imgs) + gap * (len(imgs) - 1)
        max_h = max(im.height for im in imgs)
        result = Image.new("RGB", (total_w, max_h), bg_color)
        x = 0
        for im in imgs:
            y = (max_h - im.height) // 2
            result.paste(im, (x, y))
            x += im.width + gap

    elif direction == "vertical":
        max_w = max(im.width for im in imgs)
        total_h = sum(im.height for im in imgs) + gap * (len(imgs) - 1)
        result = Image.new("RGB", (max_w, total_h), bg_color)
        y = 0
        for im in imgs:
            x = (max_w - im.width) // 2
            result.paste(im, (x, y))
            y += im.height + gap

    elif direction == "grid":
        cols = _calc_cols(len(imgs))
        rows = (len(imgs) + cols - 1) // cols

        max_w = max(im.width for im in imgs)
        max_h = max(im.height for im in imgs)
        total_w = cols * max_w + gap * (cols - 1)
        total_h = rows * max_h + gap * (rows - 1)
        result = Image.new("RGB", (total_w, total_h), bg_color)

        for idx, im in enumerate(imgs):
            r, c = divmod(idx, cols)
            x = c * (max_w + gap) + (max_w - im.width) // 2
            y = r * (max_h + gap) + (max_h - im.height) // 2
            result.paste(im, (x, y))
    else:
        raise ValueError(f"不支持的方向: {direction}")

    buf = io.BytesIO()
    result.save(buf, format="PNG")
    return buf.getvalue()


def _calc_cols(n):
    """计算宫格列数"""
    if n <= 1:
        return 1
    if n <= 4:
        return 2
    if n <= 9:
        return 3
    return 4