import os
import argparse
from PIL import Image
import concurrent.futures

# 确保使用的是 os.cpu_count() 来自动检测核心数
try:
    DEFAULT_WORKERS = os.cpu_count() or 1
except AttributeError:
    try:
        import multiprocessing
        DEFAULT_WORKERS = multiprocessing.cpu_count()
    except (ImportError, NotImplementedError):
        DEFAULT_WORKERS = 4

def _process_single_file(source_path, png_path, delete_source):
    """
    工作函数：处理单个文件的转换、缩放，并保存为PNG。
    (此函数与上一版完全相同)
    """
    filename = os.path.basename(source_path)
    png_filename = os.path.basename(png_path)
    
    try:
        with Image.open(source_path) as img:
            width, height = img.size
            min_side = min(width, height)
            resize_info = ""
            
            if min_side > 1080:
                scale_factor = 1080 / min_side
                new_width = int(width * scale_factor)
                new_height = int(height * scale_factor)
                img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
                resize_info = f" (已缩放至 {new_width}x{new_height})"
            
            img.save(png_path, 'PNG')
        
        if delete_source:
            os.remove(source_path)
            return "success_del", filename, f"{png_filename}{resize_info}"
        else:
            return "success", filename, f"{png_filename}{resize_info}"
            
    except Exception as e:
        return "fail", filename, str(e)

def process_images(folder_path, delete_source=False, max_workers=DEFAULT_WORKERS):
    """
    批量将指定文件夹 (及其所有子文件夹) 内的 .bmp, .jpg, .jpeg 图像
    转换为 .png, 并按比例缩放 (最小边<=1080px)。
    """
    
    if not os.path.isdir(folder_path):
        print(f"错误：文件夹 '{folder_path}' 不存在或不是一个有效的目录。")
        return

    print(f"--- 开始递归处理文件夹: {folder_path} ---")
    print(f"--- 使用最多 {max_workers} 个进程并行处理 ---")
    if delete_source:
        print("警告：已启用源文件删除模式。")

    # --- 1. 收集所有需要处理的任务 ---
    tasks = []
    supported_extensions = ('.bmp', '.jpg', '.jpeg')
    
    # <--- 更改：从 os.listdir() 切换到 os.walk() 以支持递归
    print("--- 正在扫描所有子文件夹...")
    for dirpath, dirnames, filenames in os.walk(folder_path):
        for filename in filenames:
            if filename.lower().endswith(supported_extensions):
                
                # 构建完整的文件路径
                source_path = os.path.join(dirpath, filename)
                
                # 构建PNG输出路径 (保持在同一个子文件夹内)
                base_name = os.path.splitext(filename)[0]
                png_path = os.path.join(dirpath, base_name + '.png')
                
                # 避免重复转换
                if os.path.exists(png_path):
                    # 打印相对路径，使其更清晰
                    relative_path = os.path.relpath(png_path, folder_path)
                    print(f"  [跳过] {relative_path} (目标文件已存在)")
                    continue

                tasks.append((source_path, png_path, delete_source))
    # <--- 更改结束 ---

    if not tasks:
        print(f"未在文件夹中找到新的 {supported_extensions} 文件。")
        return

    print(f"--- 发现 {len(tasks)} 个新图像文件，开始转换... ---")

    converted_count = 0
    failed_count = 0

    # --- 2. 使用 ProcessPoolExecutor 执行任务 (此部分不变) ---
    with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
        future_to_source = {
            executor.submit(_process_single_file, source, png, delete): source
            for source, png, delete in tasks
        }

        # --- 3. 实时获取已完成的结果 (此部分不变) ---
        for future in concurrent.futures.as_completed(future_to_source):
            source_path_orig = future_to_source[future]
            
            try:
                status, name, result = future.result()
                
                # 获取相对路径以便于阅读
                relative_dir = os.path.relpath(os.path.dirname(source_path_orig), folder_path)
                # 如果是根目录，relative_dir 会是 "."，我们将其替换为空字符串
                if relative_dir == ".":
                    log_name = name
                else:
                    log_name = os.path.join(relative_dir, name)

                if status == "success":
                    print(f"  [成功] {log_name} -> {result}")
                    converted_count += 1
                elif status == "success_del":
                    print(f"  [成功] {log_name} -> {result} (并已删除源文件)")
                    converted_count += 1
                elif status == "fail":
                    print(f"  [失败] 转换 {log_name} 时出错: {result} (源文件未删除)")
                    failed_count += 1
                    
            except Exception as e:
                print(f"  [严重失败] 处理 {os.path.basename(source_path_orig)} 时进程出错: {e}")
                failed_count += 1

    print("--- 处理完毕 ---")
    if converted_count > 0:
        print(f"成功转换 {converted_count} 个文件。")
    if failed_count > 0:
        print(f"失败 {failed_count} 个文件。")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="将指定文件夹(及其所有子文件夹)中的所有 .bmp, .jpg, .jpeg 文件批量转换为 .png 文件 (并行加速)，并按比例缩放 (最小边<=1080px)。"
    )
    
    parser.add_argument(
        "folder",
        type=str,
        help="包含 .bmp, .jpg, .jpeg 图像的 (根) 目标文件夹路径 (例如 ./ABC)"
    )
    
    parser.add_argument(
        "-d", "--delete-source",
        action="store_true",
        help="在成功转换为 .png 后，删除原始的 .bmp/.jpg/.jpeg 文件。"
    )
    
    parser.add_argument(
        "-w", "--workers",
        type=int,
        default=DEFAULT_WORKERS,
        help=f"指定用于转换的工作进程数 (默认: {DEFAULT_WORKERS}, 即本机CPU核心数)"
    )
    
    args = parser.parse_args()
    
    process_images(args.folder, args.delete_source, args.workers)