# coding=utf-8
import os
from datetime import datetime
from PIL import Image


def merge_images(image_folder, output_file, file_name, output_type, n=1):
    # 获取所有图像文件的列表
    image_files = [f for f in os.listdir(image_folder) if f.endswith('.png')]

    # 计算每个小图像的大小和大图像的大小
    image_count = len(image_files)
    if image_count == 0:
        print('No image files found in the directory:', image_folder)
        return

    # 将所有小图像粘贴到新图像的正确位置
    target_img = {}
    new_img_size1 = 0
    new_img_size0 = 0
    for i, f in enumerate(image_files):
        idx = f.strip('.png')
        target_img[int(idx)] = f
        # 计算小图像的大小以及大图像的大小
        img = Image.open(os.path.join(image_folder, image_files[i]))
        img_size1 = int(img.size[1] / n)
        if i == 0:
            new_img_size0 = int(img.size[0] / n)
        new_img_size1 += img_size1

    # 创建一个新的大图像
    new_img = Image.new('RGB', (new_img_size0, new_img_size1), 'white')

    # 将所有小图像粘贴到新图像的正确位置
    now_height = 0
    for i in range(image_count):
        img = Image.open(os.path.join(image_folder, target_img[i]))
        # 纵向拼接图片
        new_img.paste(img, (0, now_height))
        # 当前粘贴高度自增已拼接的图片高度
        now_height += img.size[1]

    # 默认两种类型都生成
    if output_type == "":
        new_img.save(r'{}/{}.png'.format(output_file, file_name))
        new_img.save(r'{}/{}.pdf'.format(output_file, file_name), 'PDF', resolution=100.0, save_all=True)
    elif output_type == "img":
        new_img.save(r'{}/{}.png'.format(output_file, file_name))
    elif output_type == "pdf":
        new_img.save(r'{}/{}.pdf'.format(output_file, file_name), 'PDF', resolution=100.0, save_all=True)


# 调用截图函数
if __name__ == "__main__":
    # 用法示例
    image_folder = './imgs'
    output_file = './screenshot'
    file_name = 'output1'
    merge_images(image_folder, output_file, file_name=file_name)