import os
import random
import shutil
import time

# ================= 🔧 配置区域 =================
# 请确保这个路径和你截图里的路径一模一样
root_path = r"D:\data\20251210\20251210"

# 划分比例 8:1:1
split_ratios = {"train": 0.8, "val": 0.1, "test": 0.1}

# 后缀配置
valid_image_exts = ['.jpg', '.jpeg', '.png', '.bmp']
valid_label_exts = ['.txt', '.xml', '.json']
# ===============================================

def main():
    print("==========================================")
    print("      正在启动 8:1:1 随机划分程序")
    print("==========================================\n")

    images_root = os.path.join(root_path, "images")
    labels_root = os.path.join(root_path, "labels")

    # 1. 检查根目录
    if not os.path.exists(images_root):
        print(f"❌ 错误：找不到 images 文件夹！\n   试图寻找路径: {images_root}")
        return
    if not os.path.exists(labels_root):
        print(f"❌ 错误：找不到 labels 文件夹！\n   试图寻找路径: {labels_root}")
        return

    print(f"✅ 路径检查通过: {root_path}")
    
    # 2. 收集文件
    all_pairs = []
    # 扫描 images 下的所有文件夹（包括 train, test, val 或者其他）
    # 只要是在 images 目录下的子文件夹，都会被扫描
    sub_dirs = [d for d in os.listdir(images_root) if os.path.isdir(os.path.join(images_root, d))]
    
    # 如果 images 下没有子文件夹，可能是图片直接放在了 images 根目录下？
    # 为了兼容，如果下面没有文件夹，就扫描 images 本身
    if not sub_dirs: 
        sub_dirs = ["."] # 代表当前目录
        print("⚠️ 提示：images 下没有子文件夹，将扫描 images 根目录...")

    print(f"📂 正在扫描以下文件夹: {sub_dirs}")

    for sub in sub_dirs:
        # 处理路径：如果是 "." 则不拼接子目录
        sub_img_dir = images_root if sub == "." else os.path.join(images_root, sub)
        sub_lbl_dir = labels_root if sub == "." else os.path.join(labels_root, sub)
        
        if not os.path.exists(sub_lbl_dir):
            # 如果对应的 label 文件夹不存在，跳过
            continue

        files = os.listdir(sub_img_dir)
        count_folder = 0
        for f in files:
            stem, ext = os.path.splitext(f)
            if ext.lower() in valid_image_exts:
                img_path = os.path.join(sub_img_dir, f)
                
                # 找标签
                lbl_path = None
                for lbl_ext in valid_label_exts:
                    potential_lbl = os.path.join(sub_lbl_dir, stem + lbl_ext)
                    if os.path.exists(potential_lbl):
                        lbl_path = potential_lbl
                        break
                
                if lbl_path:
                    all_pairs.append({'img': img_path, 'lbl': lbl_path})
                    count_folder += 1
        
        print(f"   -> 在 [{sub}] 中找到 {count_folder} 对数据")

    total = len(all_pairs)
    print(f"\n📦 总共收集到: {total} 组数据")
    
    if total == 0:
        print("❌ 没有找到任何匹配的图片和标签，请检查文件名是否对应（例如 001.jpg 是否有 001.txt）。")
        return

    # 3. 打乱与计算
    print("🎲 正在打乱顺序...")
    random.shuffle(all_pairs)
    
    n_train = int(total * split_ratios["train"])
    n_val = int(total * split_ratios["val"])
    n_test = total - n_train - n_val
    
    print(f"📊 划分数量 -> Train: {n_train}, Val: {n_val}, Test: {n_test}")

    split_data = {
        "train": all_pairs[:n_train],
        "val": all_pairs[n_train : n_train + n_val],
        "test": all_pairs[n_train + n_val :]
    }

    # 4. 移动文件
    print("🚚 开始移动文件...")
    
    for split_name, items in split_data.items():
        # 目标目录
        dst_img_dir = os.path.join(images_root, split_name)
        dst_lbl_dir = os.path.join(labels_root, split_name)
        
        os.makedirs(dst_img_dir, exist_ok=True)
        os.makedirs(dst_lbl_dir, exist_ok=True)
        
        processed = 0
        for item in items:
            src_img, src_lbl = item['img'], item['lbl']
            
            # 只有当源路径不在目标路径时才移动
            if os.path.dirname(src_img) != dst_img_dir:
                try:
                    shutil.move(src_img, os.path.join(dst_img_dir, os.path.basename(src_img)))
                    shutil.move(src_lbl, os.path.join(dst_lbl_dir, os.path.basename(src_lbl)))
                except Exception as e:
                    print(f"   [Error] 移动失败: {os.path.basename(src_img)} -> {e}")
            
            processed += 1
            # 每处理 500 个打印一次进度，防止看着像死机
            if processed % 500 == 0:
                print(f"   [{split_name}] 已处理 {processed} / {len(items)}")

        print(f"   ✅ {split_name} 完成。")

    print("\n🎉🎉🎉 全部处理完毕！ 🎉🎉🎉")

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"\n❌ 发生严重错误: {e}")
    
    # 这里的 input 是为了防止窗口直接关闭
    print("\n--------------------------------")
    input("程序运行结束，请按回车键关闭窗口...")