import os

# ================= 核心配置 =================
# 根据你的截图，你的根目录应该是这一层（包含 images 和 labels 的上一级）
# 请务必核对这个路径是否正确！
root_path = r"D:\data\20251210\20251210"

img_folder_name = "images"
lbl_folder_name = "labels"
sub_folders = ["test", "train", "val"]
global_index = 0
digit_padding = 6
# 支持的后缀（已忽略大小写）
valid_image_exts = ['.jpg', '.jpeg', '.png', '.bmp']
valid_label_exts = ['.txt', '.xml', '.json']
# ===========================================

def debug_sync_rename():
    global global_index
    
    print("----------- 开始检查路径 -----------")
    if not os.path.exists(root_path):
        print(f"❌ 错误：根目录不存在 -> {root_path}")
        print("请检查路径中是否有空格、拼写错误，或者盘符不对。")
        input("按回车键退出...")
        return
    else:
        print(f"✅ 根目录存在: {root_path}")

    images_root = os.path.join(root_path, img_folder_name)
    labels_root = os.path.join(root_path, lbl_folder_name)

    if not os.path.exists(images_root):
        print(f"❌ 错误：找不到 images 文件夹 -> {images_root}")
        input("按回车键退出...")
        return
    
    print(f"✅ images 目录存在: {images_root}")
    print(f"✅ labels 目录 (期望路径): {labels_root}")
    print("------------------------------------")

    # 遍历子文件夹
    for sub in sub_folders:
        curr_img_dir = os.path.join(images_root, sub)
        curr_lbl_dir = os.path.join(labels_root, sub)

        print(f"\n正在扫描子文件夹: [{sub}]")
        
        if not os.path.exists(curr_img_dir):
            print(f"⚠️ 跳过：找不到图片子目录 -> {curr_img_dir}")
            continue

        # 获取文件
        try:
            files = os.listdir(curr_img_dir)
        except Exception as e:
            print(f"❌ 读取文件夹出错: {e}")
            continue

        # 筛选图片
        img_files = sorted([f for f in files if os.path.splitext(f)[1].lower() in valid_image_exts])
        
        if len(img_files) == 0:
            print(f"⚠️ 在 {sub} 中没有找到支持的图片文件！")
            print(f"   该文件夹下的前5个文件是: {files[:5]}")
            continue
        
        print(f"   -> 发现 {len(img_files)} 张图片，准备处理...")

        # 开始重命名循环
        count = 0
        for old_img_name in img_files:
            stem, img_ext = os.path.splitext(old_img_name)
            
            # 构造新名字
            new_stem = f"{global_index:0{digit_padding}d}"
            
            # 1. 改名图片
            old_img_path = os.path.join(curr_img_dir, old_img_name)
            new_img_name = f"{new_stem}{img_ext}"
            new_img_path = os.path.join(curr_img_dir, new_img_name)
            
            # 2. 改名标签
            label_renamed = False
            for lbl_ext in valid_label_exts:
                old_lbl_name = f"{stem}{lbl_ext}"
                old_lbl_path = os.path.join(curr_lbl_dir, old_lbl_name)
                
                if os.path.exists(old_lbl_path):
                    new_lbl_name = f"{new_stem}{lbl_ext}"
                    new_lbl_path = os.path.join(curr_lbl_dir, new_lbl_name)
                    
                    if old_lbl_path != new_lbl_path:
                        os.rename(old_lbl_path, new_lbl_path)
                    label_renamed = True
                    break # 找到对应标签就不找其他后缀了
            
            # 执行图片改名
            if old_img_path != new_img_path:
                os.rename(old_img_path, new_img_path)

            global_index += 1
            count += 1
            
            # 为了避免刷屏，只打印前3个和最后一个的处理记录
            if count <= 3:
                status = "同步修改成功" if label_renamed else "只改了图片(无标签)"
                print(f"   [{status}] {old_img_name} -> {new_img_name}")

        print(f"   -> {sub} 文件夹处理完毕。")

    print(f"\n🎉 所有操作结束！最大序号停留在: {global_index - 1}")
    input("按回车键退出...")

if __name__ == "__main__":
    debug_sync_rename()