import os
import shutil
import argparse

def delete_json_files(folder_path):
    """删除指定文件夹中的所有JSON文件"""
    for filename in os.listdir(folder_path):
        if filename.endswith('.json'):
            file_path = os.path.join(folder_path, filename)
            try:
                os.remove(file_path)
                print(f"已删除JSON文件: {file_path}")
            except Exception as e:
                print(f"删除 {file_path} 时出错: {str(e)}")

def merge_folders(input_path, output_dir):
    """合并三个文件夹的内容到输出目录"""
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 定义要处理的子文件夹
    subfolders = ['train', 'test', 'val']
    
    # 检查输入目录结构有效性
    for sub in subfolders:
        if not os.path.exists(os.path.join(input_path, sub)):
            raise ValueError(f"输入路径中缺少 {sub} 文件夹")

    # 处理每个子文件夹
    for sub in subfolders:
        sub_path = os.path.join(input_path, sub)
        
        # 删除JSON文件
        delete_json_files(sub_path)
        
        # 复制所有文件到目标目录
        for filename in os.listdir(sub_path):
            src = os.path.join(sub_path, filename)
            dst = os.path.join(output_dir, filename)
            
            # 处理重复文件名（保留后复制到的文件）
            if os.path.exists(dst):
                print(f"警告: {filename} 已存在，将被覆盖")
                
            shutil.copy2(src, dst)
            print(f"已复制: {src} -> {dst}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="合并训练/测试/验证集并删除JSON文件")
    parser.add_argument('-i','--input_path', required=True, help='包含train/test/val文件夹的输入路径')
    parser.add_argument('-o','--output_dir', required=True, help='合并后的输出目录路径')
    
    args = parser.parse_args()
    
    try:
        print(f"开始处理，输入路径: {args.input_path}")
        merge_folders(args.input_path, args.output_dir)
        print("\n操作完成！合并后的文件位于:", os.path.abspath(args.output_dir))
    except Exception as e:
        print(f"\n发生错误: {str(e)}")