# -*- coding: utf-8 -*-

import os
import re
import sys
import subprocess
import shutil
import tempfile
import zipfile

# --- 配置区 ---
IMAGE_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.bmp', '.gif', '.tiff')
# --- 配置区结束 ---

def check_ffmpeg():
    """检查系统中是否安装了 ffmpeg"""
    if shutil.which("ffmpeg") is None:
        print("错误：找不到 ffmpeg。")
        print("请确保你已经安装了 ffmpeg，并将其添加到了系统的 PATH 环境变量中。")
        print("安装指南: https://ffmpeg.org/download.html")
        return False
    return True

def extract_number_from_filename(filename):
    """从文件名中提取最后一个遇到的数字序列用于排序。"""
    numbers = re.findall(r'\d+', filename)
    if numbers:
        return int(numbers[-1])
    print(f"警告：文件名 '{filename}' 中未找到数字，将按0进行排序。")
    return 0

def main():
    """主执行函数"""
    print("--- 图片序列转 WebP 并压缩为 ZIP 脚本 (支持自定义起始编号) ---")

    # 1. 检查命令行参数 (现在支持3个或4个参数)
    if not (3 <= len(sys.argv) <= 4):
        print("\n使用方法:")
        print(f"  python {sys.argv[0]} <源图片文件夹> <输出的ZIP文件名> [<起始编号>]")
        print("\n示例 (从0开始):")
        print(f"  python {sys.argv[0]} ./my_images my_archive.zip")
        print("\n示例 (从40开始):")
        print(f"  python {sys.argv[0]} ./my_images my_archive.zip 40")
        sys.exit(1)

    source_dir = sys.argv[1]
    output_zip_name = sys.argv[2]
    
    # 初始化起始编号为默认值 0
    start_index = 0
    
    # 如果提供了第三个参数，则用它作为起始编号
    if len(sys.argv) == 4:
        try:
            start_index = int(sys.argv[3])
            if start_index < 0:
                print("错误：起始编号不能为负数。")
                sys.exit(1)
        except ValueError:
            print(f"错误：提供的起始编号 '{sys.argv[3]}' 不是一个有效的整数。")
            sys.exit(1)

    if not output_zip_name.lower().endswith('.zip'):
        output_zip_name += '.zip'

    # 2. 验证输入
    if not check_ffmpeg():
        sys.exit(1)
        
    if not os.path.isdir(source_dir):
        print(f"错误：源文件夹 '{source_dir}' 不存在或不是一个目录。")
        sys.exit(1)

    # 3. 查找并排序图片文件
    print(f"\n[1/5] 正在扫描文件夹: {source_dir}")
    try:
        all_files = os.listdir(source_dir)
        image_files = [f for f in all_files if f.lower().endswith(IMAGE_EXTENSIONS)]
    except OSError as e:
        print(f"错误：无法读取文件夹 '{source_dir}': {e}")
        sys.exit(1)

    if not image_files:
        print("指定的文件夹中未找到任何支持的图片文件。")
        sys.exit(0)

    sorted_images = sorted(image_files, key=extract_number_from_filename)
    
    print("找到并排序后的图片文件：")
    for img in sorted_images:
        print(f"  - {img}")

    # 使用临时目录来存放生成的 webp 文件
    with tempfile.TemporaryDirectory() as temp_dir:
        print(f"\n[2/5] 开始转换图片 (起始编号: {start_index})...")
        total_files = len(sorted_images)
        
        # 4. 依次使用 ffmpeg 进行编码
        for index, image_name in enumerate(sorted_images):
            # 计算最终的输出文件名编号
            output_number = start_index + index
            output_filename = f"{output_number}.webp"
            
            input_path = os.path.join(source_dir, image_name)
            output_path = os.path.join(temp_dir, output_filename)
            
            command = ["ffmpeg", "-i", input_path, "-y", output_path]
            
            print(f"  [{index + 1}/{total_files}] 正在转换: {image_name} -> {output_filename}")
            
            try:
                subprocess.run(
                    command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, encoding='utf-8'
                )
            except subprocess.CalledProcessError as e:
                print(f"    错误: 转换 '{image_name}' 时失败。\n    FFmpeg 错误信息:\n{e.stderr}")
            except Exception as e:
                print(f"    发生未知错误: {e}")

        # 5. 将临时目录中的 WebP 文件压缩成 ZIP
        print(f"\n[3/5] 正在将 WebP 文件压缩到: {output_zip_name}")
        
        try:
            with zipfile.ZipFile(output_zip_name, 'w', zipfile.ZIP_DEFLATED) as zipf:
                webp_files_to_zip = sorted(os.listdir(temp_dir), key=lambda f: int(f.split('.')[0]))
                
                for webp_file in webp_files_to_zip:
                    file_path_in_temp = os.path.join(temp_dir, webp_file)
                    zipf.write(file_path_in_temp, arcname=webp_file)
                    print(f"  - 已添加 {webp_file} 到 ZIP 包")

        except Exception as e:
            print(f"错误：创建 ZIP 文件时失败: {e}")
            sys.exit(1)

    print("\n[4/5] 临时文件已自动清理。")
    print(f"\n[5/5] --- 所有任务完成！输出文件已保存为 '{output_zip_name}' ---")


if __name__ == "__main__":
    main()