import os
import shutil
import math

# ================= 配置区域 =================
# 【全局变量】在这里修改你要分成的份数
N_FOLDERS = 4

# 源文件夹名称 (相对于当前脚本运行目录)
SOURCE_DIR_NAME = "camera02_save"

# 支持的图片扩展名
VALID_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff')


# ===========================================

def split_images_to_subfolders():
    # 1. 获取当前脚本所在的绝对路径
    base_dir = os.getcwd()
    source_path = os.path.join(base_dir, SOURCE_DIR_NAME)

    # 2. 检查源文件夹是否存在
    if not os.path.exists(source_path):
        print(f"❌ 错误: 找不到文件夹 '{source_path}'")
        print("请确保脚本在包含 'camera02_save' 的目录下运行。")
        return

    # 3. 获取所有图片文件列表
    all_files = os.listdir(source_path)
    image_files = [
        f for f in all_files
        if f.lower().endswith(VALID_EXTENSIONS) and os.path.isfile(os.path.join(source_path, f))
    ]

    # 按文件名排序，保证分配顺序一致（可选）
    image_files.sort()

    total_count = len(image_files)

    if total_count == 0:
        print(f"警告: '{source_path}' 中没有找到任何图片文件。")
        return

    print(f"发现图片总数: {total_count}")
    print(f"计划分成: {N_FOLDERS} 份")

    # 4. 计算分配逻辑
    # 使用轮询方式分配，确保每个文件夹数量尽可能平均
    # 例如：10张图分3份 -> 4, 3, 3

    created_folders = []

    for index, filename in enumerate(image_files):
        # 计算该图片应该去哪个子文件夹 (0 到 N_FOLDERS-1)
        folder_index = index % N_FOLDERS
        subfolder_name = f"part_{folder_index}"

        # 构建子文件夹完整路径
        dest_folder_path = os.path.join(source_path, subfolder_name)

        # 如果文件夹不存在，则创建
        if not os.path.exists(dest_folder_path):
            os.makedirs(dest_folder_path)
            created_folders.append(subfolder_name)
            print(f"✅ 创建子文件夹: {subfolder_name}")

        # 构建源文件和目标文件路径
        src_file = os.path.join(source_path, filename)
        dst_file = os.path.join(dest_folder_path, filename)

        # 移动文件 (剪切)
        try:
            shutil.move(src_file, dst_file)
        except Exception as e:
            print(f"❌ 移动失败 {filename}: {e}")

    # 5. 统计结果
    print("\n" + "=" * 30)
    print("分配完成！统计如下：")
    print("=" * 30)

    # 重新扫描统计每个文件夹的数量
    final_counts = {}
    for i in range(N_FOLDERS):
        fname = f"part_{i}"
        fpath = os.path.join(source_path, fname)
        if os.path.exists(fpath):
            count = len([f for f in os.listdir(fpath) if f.lower().endswith(VALID_EXTENSIONS)])
            final_counts[fname] = count
            print(f"{fname}: {count} 张图片")

    print("=" * 30)
    print(f"总计移动: {sum(final_counts.values())} 张图片")


if __name__ == "__main__":
    # 提示用户确认
    print(f"当前工作目录: {os.getcwd()}")
    confirm = input(f"即将把 '{SOURCE_DIR_NAME}' 中的图片均分为 {N_FOLDERS} 份。\n是否继续？(y/n): ")

    if confirm.lower() == 'y':
        split_images_to_subfolders()
    else:
        print("操作已取消。")