import shutil
import sys
import zipfile
import os
from PIL import Image
from fpdf import FPDF


def read_config(config_path):
    # 默认配置
    default_config = {
        'imagesPerRow': 3,
        'rowsPerPage': 2,
        'imageWidth': 60,
        'imageHeight': 130,
        'xStart': 10,
        'yStart': 20
    }

    # 如果配置文件不存在，返回默认配置
    if not os.path.exists(config_path):
        return default_config

    # 读取配置文件
    config = default_config.copy()
    with open(config_path, 'r') as file:
        for line in file:
            name, value = line.split('=')
            config[name.strip()] = int(value.strip())
    return config


def collect_images(src_dir, tmp_dir):
    os.makedirs(tmp_dir, exist_ok=True)
    # 支持更多图像格式
    supported_formats = ('.jpg', '.jpeg', '.png', '.tiff', '.bmp', '.gif')

    for root, _, files in os.walk(src_dir):
        for file_name in files:
            if file_name.endswith(supported_formats):
                src_file_path = os.path.join(root, file_name)
                dst_file_path = os.path.join(tmp_dir, file_name)
                shutil.copy(src_file_path, dst_file_path)
            elif file_name.endswith('.zip'):
                zip_file_path = os.path.join(root, file_name)
                extract_path = os.path.join(tmp_dir, os.path.splitext(file_name)[0])
                os.makedirs(extract_path, exist_ok=True)
                with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
                    zip_ref.extractall(extract_path)
                for extracted_file in os.listdir(extract_path):
                    extracted_file_path = os.path.join(extract_path, extracted_file)
                    if extracted_file.endswith(supported_formats):
                        dst_file_path = os.path.join(tmp_dir, extracted_file)
                        shutil.move(extracted_file_path, dst_file_path)
                shutil.rmtree(extract_path, ignore_errors=True)


def process_images_to_pdf(tmp_dir, output_pdf_path, config):
    try:
        # 支持更多图像格式
        supported_formats = ('.jpg', '.jpeg', '.png', '.tiff', '.bmp', '.gif')

        # 收集所有支持格式的图像文件
        extracted_files = [f for f in os.listdir(tmp_dir) if f.endswith(supported_formats)]

        # Sort the files in descending order based on the timestamp in the filename
        sorted_files_descending = sorted(extracted_files, key=lambda x: int(x.split('.')[0]), reverse=True)

        # Create a PDF document with images arranged in rows, three images per row, six images per page,
        # proportionally scaled
        pdf = FPDF()

        # 读取布局参数
        images_per_row = config['imagesPerRow']
        rows_per_page = config['rowsPerPage']
        max_images_per_page = images_per_row * rows_per_page
        image_width = config['imageWidth']  # 每张图片的宽度比例
        image_height = config['imageHeight']  # 每张图片的最大高度以适应两行

        # 在 PDF 中按比例缩放添加图片
        x_start = config['xStart']
        y_start = config['yStart']
        current_x = x_start
        current_y = y_start

        # Function to convert images to supported format
        def convert_image_to_jpeg(image_path):
            img = Image.open(image_path)
            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")
            jpeg_path = image_path + ".jpg"
            img.save(jpeg_path, "JPEG")
            return jpeg_path

        # Function to get scaled dimensions
        def get_scaled_dimensions(image_path, max_width, max_height):
            img = Image.open(image_path)
            width, height = img.size
            ratio = min(max_width / width, max_height / height)
            return width * ratio, height * ratio

        for idx, image_file in enumerate(sorted_files_descending):
            if idx % max_images_per_page == 0:
                pdf.add_page()
                current_x = x_start
                current_y = y_start

            image_path = os.path.join(tmp_dir, image_file)
            jpeg_path = convert_image_to_jpeg(image_path)
            width, height = get_scaled_dimensions(jpeg_path, image_width, image_height)
            pdf.image(jpeg_path, x=current_x, y=current_y, w=width, h=height)

            # Update x position for next image
            current_x += image_width + 5  # Add some margin between images

            # If we have placed three images in the current row, move to the next row
            if (idx + 1) % images_per_row == 0:
                current_x = x_start
                current_y += image_height + 10  # Add some margin between rows

        # Save the PDF to a file
        pdf.output(output_pdf_path)
        print(f'PDF saved to: {output_pdf_path}')
    finally:
        # Clean up tmp directory if needed
        shutil.rmtree(tmp_dir, ignore_errors=True)


def main():
    current_dir = os.getcwd()
    config_path = os.path.join(current_dir, 'config.txt')
    tmp_dir = os.path.join(current_dir, 'tmp')
    output_dir_path = os.path.join(current_dir, 'out')
    origin_dir_path = os.path.join(current_dir, 'orgin')

    # 读取配置文件
    config = read_config(config_path)

    # Ensure output directory exists
    os.makedirs(output_dir_path, exist_ok=True)

    # Collect images from 'orgin' directory
    collect_images(origin_dir_path, tmp_dir)

    # Generate PDF from collected images
    output_pdf_path = os.path.join(output_dir_path, 'output.pdf')
    process_images_to_pdf(tmp_dir, output_pdf_path, config)


if __name__ == "__main__":
    main()