import os
import sys

def find_specific_epochs(root_folder):
    """
    遍历指定文件夹，查找所有 epoch_XXX.pth 文件，
    其中 XXX 为整数且不能被 10 整除，并返回它们的绝对路径。

    :param root_folder: 要搜索的根文件夹路径
    :return: 一个包含符合条件文件绝对路径的列表
    """
    matching_files = []
    
    for dirpath, _, filenames in os.walk(root_folder):
        for filename in filenames:
            if filename.startswith('epoch_') and filename.endswith('.pth'):
                try:
                    number_str = filename[len('epoch_'):-len('.pth')]
                    if number_str.isdigit():
                        epoch_number = int(number_str)
                        if epoch_number % 10 != 0:
                            full_path = os.path.join(dirpath, filename)
                            matching_files.append(os.path.abspath(full_path))
                except (ValueError, IndexError):
                    continue
                    
    return matching_files

if __name__ == "__main__":
    # 1. 获取当前脚本文件所在的绝对目录
    script_dir = os.path.dirname(os.path.abspath(__file__))

    # 2. 以脚本目录为基准，构建目标路径
    # V1
    target_directory = os.path.abspath(os.path.join(script_dir, '../../Hardisk'))
    # V2
    # target_directory = os.path.abspath(os.path.join(script_dir, '../../DataSet_Public_outputs'))

    # 1. 查找符合条件的文件
    found_paths = find_specific_epochs(target_directory)
    
    if not found_paths:
        print(f"在 '{os.path.abspath(target_directory)}' 及其子目录中没有找到符合条件的文件。")
        sys.exit(0)

    # 2. 列出所有找到的文件，并请求用户确认
    print("找到了以下符合条件的文件：")
    for path in found_paths:
        print(path)
    
    print("\n" + "="*50)
    print("警告：接下来的操作将永久删除以上列出的所有文件！")
    print("="*50 + "\n")
    
    # 3. 获取用户输入
    try:
        # 使用 strip() 去除首尾空格，使用 lower() 转换为小写
        confirm = input("您确定要删除这 {} 个文件吗？ (请输入 'yes' 进行确认，输入其他任何内容则取消): ".format(len(found_paths)))
    except KeyboardInterrupt:
        print("\n\n操作被用户中断。")
        sys.exit(1)

    # 4. 根据用户的输入执行操作
    if confirm.lower().strip() == 'yes':
        print("\n正在开始删除文件...")
        deleted_count = 0
        error_count = 0
        for path in found_paths:
            try:
                os.remove(path)
                print(f"已删除: {path}")
                deleted_count += 1
            except OSError as e:
                print(f"删除失败: {path} (原因: {e})")
                error_count += 1
        print(f"\n操作完成。成功删除 {deleted_count} 个文件，{error_count} 个文件删除失败。")
    else:
        print("\n操作已取消，没有文件被删除。")