#!/usr/bin/env python3
import argparse
import re
from pathlib import Path

from PIL import Image

IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tga", ".webp"}


def positive_int(text):
    value = int(text)
    if value < 1:
        raise argparse.ArgumentTypeError("必须是正整数")
    return value


def natural_key(path):
    return [int(part) if part.isdigit() else part.lower() for part in re.split(r"(\d+)", path.name)]


def parse_args():
    parser = argparse.ArgumentParser(description="把文件夹内的图片按行列拼接成一张 PNG 图集，空单元格完全透明。")
    parser.add_argument("input", type=Path, help="输入图片文件夹（不递归子文件夹）")
    parser.add_argument("-r", "--rows", type=positive_int, required=True, help="行数")
    parser.add_argument("-c", "--cols", type=positive_int, required=True, help="列数")
    parser.add_argument("-s", "--cell-size", type=positive_int, required=True, help="单元格边长（像素）")
    parser.add_argument("-o", "--output", type=Path, help="输出 PNG 路径，默认为输入文件夹同级目录下的 <文件夹名>_atlas.png")
    return parser.parse_args()


def collect_images(folder):
    files = [path for path in folder.iterdir() if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS]
    files.sort(key=natural_key)
    return files


def build_atlas(files, rows, cols, cell_size):
    atlas = Image.new("RGBA", (cols * cell_size, rows * cell_size), (0, 0, 0, 0))
    for index, path in enumerate(files):
        row, col = divmod(index, cols)
        try:
            with Image.open(path) as image:
                cell = image.convert("RGBA").resize((cell_size, cell_size), Image.LANCZOS)
        except OSError as error:
            raise SystemExit(f"错误：无法读取图片 {path}：{error}")
        atlas.paste(cell, (col * cell_size, row * cell_size))
    return atlas


def main():
    args = parse_args()
    folder = args.input
    if not folder.is_dir():
        raise SystemExit(f"错误：{folder} 不是有效的文件夹")

    files = collect_images(folder)
    if not files:
        raise SystemExit(f"错误：{folder} 中没有可用的图片文件")

    capacity = args.rows * args.cols
    if len(files) > capacity:
        raise SystemExit(f"错误：图片数量 {len(files)} 超过单元格数量 {capacity}（{args.rows} 行 x {args.cols} 列）")

    output = args.output if args.output else folder.parent / f"{folder.name}_atlas.png"
    output.parent.mkdir(parents=True, exist_ok=True)

    atlas = build_atlas(files, args.rows, args.cols, args.cell_size)
    atlas.save(output, format="PNG")
    print(f"已生成 {output}（{atlas.width}x{atlas.height}，共 {len(files)} 张图片）")


if __name__ == "__main__":
    main()