import os
import shutil
from pathlib import Path

def split_images_into_five(source_dir, target_root, move_instead_of_copy=True):
    """
    将 source_dir 中的图片均分为 5 份，保存到 target_root/1 ~ 5/

    :param source_dir: 源图片文件夹路径
    :param target_root: 目标根目录（会创建 1~5 子文件夹）
    :param move_instead_of_copy: True 表示移动，False 表示复制（默认为 True 移动）
    """
    source = Path(source_dir)
    target_root = Path(target_root)

    # 支持的图片扩展名
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif', '.webp'}

    # 获取所有图片文件（按名称排序，保证可重复性）
    all_images = sorted([
        f for f in source.iterdir()
        if f.is_file() and f.suffix.lower() in image_extensions
    ])

    if not all_images:
        print("❌ 源文件夹中没有找到图片！")
        return

    total = len(all_images)
    print(f"📁 共找到 {total} 张图片")

    # 均分为 5 份
    chunk_size = (total + 4) // 5  # 向上取整，确保覆盖所有图片

    # 创建 1~5 文件夹并分发图片
    for i in range(5):
        start_idx = i * chunk_size
        end_idx = min((i + 1) * chunk_size, total)
        if start_idx >= total:
            break  # 防止空分片

        folder_name = str(i + 1)
        target_folder = target_root / folder_name
        target_folder.mkdir(parents=True, exist_ok=True)

        chunk_files = all_images[start_idx:end_idx]
        print(f"📦 文件夹 {folder_name}: 分配 {len(chunk_files)} 张图片")

        for img_path in chunk_files:
            dst = target_folder / img_path.name
            if move_instead_of_copy:
                shutil.move(str(img_path), str(dst))
                print(f"  📂 已移动: {img_path.name}")
            else:
                shutil.copy2(str(img_path), str(dst))
                print(f"  📂 已复制: {img_path.name}")

    print(f"\n✅ 分割完成！结果保存在: {target_root}")

if __name__ == "__main__":
    # ========================
    # ⚙️ 配置你的路径
    # ========================
    SOURCE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/分割/class3"
    TARGET_ROOT = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/分割/class3"

    # 设置为 True 则移动图片（原图会被移走），False 则复制（推荐先用 False 测试）
    MOVE_MODE = True

    split_images_into_five(SOURCE_DIR, TARGET_ROOT, move_instead_of_copy=MOVE_MODE)