import os
from reportlab.lib.pagesizes import A3
from reportlab.platypus import SimpleDocTemplate, Image
from PyPDF2 import PdfReader, PdfWriter
from PIL import Image as PILImage


def create_pdf_from_images(directory, output_pdf):
    image_extensions = {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp'}
    image_files = []

    # 遍历目录并筛选图片文件
    for root, dirs, files in os.walk(directory):
        for file in files:
            _, ext = os.path.splitext(file)
            if ext.lower() in image_extensions:
                image_files.append(os.path.join(root, file))

    # 如果没有找到任何图片，返回
    if not image_files:
        print(f"No image files found in {directory}.")
        return

    # 按文件名排序，确保图片按顺序
    image_files.sort()

    # 设置PDF文档，使用 A3 页面大小
    doc = SimpleDocTemplate(output_pdf, pagesize=A3)

    # 创建图片列表
    image_list = []
    for image_file in image_files:
        # 使用 Pillow 获取图片的尺寸
        with PILImage.open(image_file) as img:
            img_width, img_height = img.size

        # 设置目标宽度（适应A3页面宽度，最大842点）
        target_width = 800  # 设置图片宽度为最大800点，适应A3页面

        # 计算新的高度，保持原始宽高比
        target_height = int(img_height * (target_width / img_width))

        # 限制图片的最大高度，避免超出页面
        max_page_height = 1091 - 100  # 留出一些空间，A3页面高度1191点
        if target_height > max_page_height:
            target_height = max_page_height
            target_width = int(target_height * (img_width / img_height))  # 调整宽度保持比例

        # 创建图片对象并添加到图片列表
        img = Image(image_file, width=target_width, height=target_height)
        image_list.append(img)

    # 生成 PDF
    doc.build(image_list)
    print(f"PDF {output_pdf} created successfully.")


def encrypt_pdf(input_pdf, output_pdf, password):
    # 使用 PyPDF2 加密 PDF
    writer = PdfWriter()
    with open(input_pdf, "rb") as file:
        reader = PdfReader(file)
        for page in reader.pages:
            writer.add_page(page)

    with open(output_pdf, "wb") as file:
        writer.encrypt(password)
        writer.write(file)
    print(f"PDF {output_pdf} encrypted successfully.")


def generate_pdf_from_images(directory):
    # 获取当前目录下所有数字命名的文件夹
    folder_names = [folder for folder in os.listdir(directory)
                    if os.path.isdir(os.path.join(directory, folder)) and folder.isdigit()]

    # 如果没有数字命名的文件夹，提示并返回
    if not folder_names:
        print("No numeric-named folders found.")
        return

    # 循环处理每个数字命名的文件夹
    for folder_name in folder_names:
        folder_path = os.path.join(directory, folder_name)

        # 设置 PDF 输出路径，使用文件夹名称作为文件名，并存储到 ./PDF 目录下
        output_pdf = f"./PDF/{folder_name}.pdf"

        # 检查 PDF 文件是否已存在
        if os.path.exists(output_pdf):
            print(f"PDF {output_pdf} already exists. Skipping...")
            continue

        # 创建目录
        os.makedirs("./PDF", exist_ok=True)

        # 创建 PDF
        create_pdf_from_images(folder_path, output_pdf)

        # 加密 PDF
        encrypt_pdf(output_pdf, output_pdf, "4000")  # 密码设置为 4000

        # 返回PDF位置，用于发送到群里
        return os.path.abspath(output_pdf)