import os
import shutil
from pathlib import Path

def copy_files_by_id(id_list_file, source_dirs, output_dir):
    """
    根据ID列表拷贝文件到新目录，保持目录结构
    
    参数:
        id_list_file: 包含文件ID列表的文本文件
        source_dirs: 源目录列表 [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
        output_dir: 输出目录
    """
    # 读取ID列表
    with open(id_list_file, 'r') as f:
        id_list = [line.strip() for line in f.readlines() if line.strip()]
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 遍历所有源目录
    for src_dir in source_dirs:
        if not os.path.exists(src_dir):
            print(f"警告: 源目录不存在 {src_dir}")
            continue
            
        # 遍历源目录下的所有文件
        for root, _, files in os.walk(src_dir):
            for file in files:
                # 检查文件名是否包含ID列表中的任一ID
                if any(file_id in file for file_id in id_list):
                    src_path = os.path.join(root, file)
                    
                    # 计算相对路径
                    rel_path = os.path.relpath(root, src_dir)
                    dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
                    
                    # 创建目标目录并拷贝文件
                    os.makedirs(dst_dir, exist_ok=True)
                    dst_path = os.path.join(dst_dir, file)
                    shutil.copy2(src_path, dst_path)
                    print(f"已拷贝: {src_path} -> {dst_path}")

if __name__ == "__main__":
    # import argparse
    
    # parser = argparse.ArgumentParser(description='根据ID列表拷贝文件')
    # parser.add_argument('id_list', help='包含文件ID列表的文本文件路径')
    # parser.add_argument('output_dir', help='输出目录路径')
    # parser.add_argument('--source_dirs', nargs='+', 
    #                     default=['hair_template_material', 'ref_hairstyle', 'train_material', 'upload_train_imgs'],
    #                     help='源目录列表')
    
    # args = parser.parse_args()
    id_list = [1939351512506802177,
        1939351392453238785,
        1939351319849836545,
        1939351238526476289,
        1939351219396255745,
        1939339435616608258,
        1939330299344560130,
        1938848782956732417,
        1938848730150445058,
        1938848709308948482,
        1938841049238970370,
        1938492830059438081,
        1938492790033195009,
        1938490058547240961,
        1938489994072399874,
        1938489915521474561,
        1938489840707674114,
        1938489739088076801,
        1938489674713899010,
        1938489620162781185]
    output_dir = '/home/data/hair/data/hairStyle_addin/06300845'
    hair_template_material = '/home/data/hair/data/hair_template_material'
    ref_hairstyle = '/home/data/hair/data/ref_hairstyle'
    train_material = '/mnt/database2/online-server/hair-online-tj-v2/train_material'
    upload_train_imgs = '/home/data/hair/data/upload_train_imgs'
    copy_files_by_id(id_list, source_dirs, args.output_dir)

def copy_files_by_id(id_list, hair_template_material, ref_hairstyle, train_material, upload_train_imgs, output_dir):
    """
    根据ID列表从指定目录拷贝文件到输出目录
    :param id_list: 文件ID列表(字符串格式)
    :param hair_template_material: 发型模板材质目录
    :param ref_hairstyle: 参考发型目录
    :param train_material: 训练材质目录
    :param upload_train_imgs: 上传训练图片目录
    :param output_dir: 输出目录
    """
    # 将ID转换为字符串格式
    id_list = [str(id) for id in id_list]
    # 创建源目录列表
    source_dirs = [hair_template_material, ref_hairstyle, train_material, upload_train_imgs]
    for src_dir in source_dirs:
        if not os.path.exists(src_dir):
            print(f"警告: 源目录不存在 {src_dir}")
            continue
            
        # 遍历源目录下的所有文件
        for root, _, files in os.walk(src_dir):
            for file in files:
                # 检查文件名是否包含ID列表中的任一ID
                if any(file_id in file for file_id in id_list):
                    src_path = os.path.join(root, file)
                    
                    # 计算相对路径
                    rel_path = os.path.relpath(root, src_dir)
                    dst_dir = os.path.join(output_dir, os.path.basename(src_dir), rel_path)
                    
                    # 创建目标目录并拷贝文件
                    os.makedirs(dst_dir, exist_ok=True)
                    dst_path = os.path.join(dst_dir, file)
                    shutil.copy2(src_path, dst_path)
                    print(f"已拷贝: {src_path} -> {dst_path}")