import os
import re
import shutil
from tqdm import tqdm

# --- 配置参数 ---
# 1. 图片文件夹路径
image_folder = '../label/up'

# 2. 标注文件夹路径
annotation_folder = '../label/up_xml'

# 3. 新文件名的数字位数 (3 -> 000, 001, ...; 5 -> 00000, 00001, ...)
padding_zeros = 3


# --- 配置结束 ---


def rename_dataset_files(img_folder, an_folder, padding=3):
    """
    重命名图片和标注文件，l系列在前，r系列在后，按数字顺序排列。
    """
    print("开始重命名任务...")

    if not os.path.isdir(img_folder):
        print(f"错误：图片文件夹 '{img_folder}' 不存在。")
        return

    if not os.path.isdir(an_folder):
        print(f"错误：标注文件夹 '{an_folder}' 不存在。")
        return

    # 1. 获取所有图片文件名，并按 l 和 r 分组
    l_files = []
    r_files = []
    other_files = []

    for filename in os.listdir(img_folder):
        if filename.lower().endswith('.jpg'):
            base_name = os.path.splitext(filename)[0]
            match = re.match(r'([lr])(\d+)', base_name, re.IGNORECASE)
            if match:
                prefix = match.group(1).lower()
                number = int(match.group(2))
                if prefix == 'l':
                    l_files.append((number, filename))
                elif prefix == 'r':
                    r_files.append((number, filename))
            else:
                other_files.append(filename)

    # 2. 对 l 和 r 组内的文件按数字进行排序
    l_files.sort()
    r_files.sort()

    # 3. 合并文件列表，l在前，r在后
    sorted_filenames = [f[1] for f in l_files] + [f[1] for f in r_files]

    if not sorted_filenames:
        print("在指定的 l/r 命名规则下未找到任何文件。")
        if other_files:
            print(f"跳过了这些文件: {other_files}")
        return

    print(f"找到 {len(sorted_filenames)} 个符合 l/r 规则的文件准备重命名。")

    # 4. 开始重命名
    counter = 0
    with tqdm(total=len(sorted_filenames), desc="重命名文件") as pbar:
        for old_filename in sorted_filenames:
            # 构建新文件名
            new_base_name = str(counter).zfill(padding)
            new_jpg_name = f"{new_base_name}.jpg"
            new_xml_name = f"{new_base_name}.xml"

            # 构建旧文件的完整路径
            old_base_name = os.path.splitext(old_filename)[0]
            old_jpg_path = os.path.join(img_folder, f"{old_base_name}.jpg")
            old_xml_path = os.path.join(an_folder, f"{old_base_name}.xml")

            # 构建新文件的完整路径
            new_jpg_path = os.path.join(img_folder, new_jpg_name)
            new_xml_path = os.path.join(an_folder, new_xml_name)

            # 执行重命名 (使用 shutil.move 更安全)
            try:
                if os.path.exists(old_jpg_path):
                    shutil.move(old_jpg_path, new_jpg_path)
                else:
                    print(f"\n警告：找不到图片文件 {old_jpg_path}")

                if os.path.exists(old_xml_path):
                    shutil.move(old_xml_path, new_xml_path)
                else:
                    print(f"\n警告：找不到XML文件 {old_xml_path}")

                counter += 1
            except Exception as e:
                print(f"\n重命名文件 {old_filename} 时出错: {e}")

            pbar.update(1)

    print(f"\n重命名完成！共处理了 {counter} 对文件。")
    if other_files:
        print(f"注意：以下文件未被处理，因为它们不符合 l/r 命名规则: {other_files}")


if __name__ == '__main__':
    rename_dataset_files(image_folder, annotation_folder, padding=padding_zeros)