import os
import shutil

def match_pcd_meta_tiff(image_json_folder, pcd_folder, depth_tiff_folder,
                        output_pcd_folder, output_meta_folder, output_tiff_folder):
    os.makedirs(output_pcd_folder, exist_ok=True)
    os.makedirs(output_meta_folder, exist_ok=True)
    os.makedirs(output_tiff_folder, exist_ok=True)

    # PCD 前缀集合（去掉后缀 .pcd）
    pcd_prefixes = {
        os.path.splitext(f)[0] for f in os.listdir(pcd_folder) if f.lower().endswith('.pcd')
    }

    # image_json_folder 中的 jpg 和 json 文件（带 _color 后缀）
    all_files = os.listdir(image_json_folder)
    color_jpgs = [f for f in all_files if f.lower().endswith('_color.jpg')]
    jsons = [f for f in all_files if f.lower().endswith('_color.json')]

    # 生成前缀映射，去掉后缀 _color.jpg / _color.json
    jpg_prefix_map = {f[:-10]: f for f in color_jpgs}    # '_color.jpg' 长度是10
    json_prefix_map = {f[:-11]: f for f in jsons}        # '_color.json' 长度是11

    # depth_tiff_folder 中的深度图文件，文件名以 _depth.tiff 或 _depth.tif 结尾
    tiff_files = [f for f in os.listdir(depth_tiff_folder) if f.lower().endswith(('.tiff', '.tif'))]
    # 去掉 _depth.tiff 后缀得到前缀
    tiff_prefix_map = {}
    for f in tiff_files:
        if f.lower().endswith('_depth.tiff'):
            prefix = f[:-11]  # 去掉 '_depth.tiff' (11个字符)
        elif f.lower().endswith('_depth.tif'):
            prefix = f[:-10]  # 去掉 '_depth.tif' (10个字符)
        else:
            continue
        tiff_prefix_map[prefix] = f

    # 找到同时有 pcd、jpg、json、tiff 的 prefix
    matched_prefixes = []
    for prefix in pcd_prefixes:
        if (prefix in jpg_prefix_map and
            prefix in json_prefix_map and
            prefix in tiff_prefix_map):
            matched_prefixes.append(prefix)

    print(f"共发现 {len(pcd_prefixes)} 个 pcd 文件")
    print(f"共发现 {len(color_jpgs)} 个 jpg (_color.jpg) 文件")
    print(f"共发现 {len(jsons)} 个 json (_color.json) 文件")
    print(f"共发现 {len(tiff_files)} 个 tiff (_depth.tiff/.tif) 文件")
    print(f"匹配到 {len(matched_prefixes)} 组文件")

    # 复制匹配文件到对应输出文件夹
    for prefix in matched_prefixes:
        pcd_src = os.path.join(pcd_folder, prefix + '.pcd')
        jpg_src = os.path.join(image_json_folder, jpg_prefix_map[prefix])
        json_src = os.path.join(image_json_folder, json_prefix_map[prefix])
        tiff_src = os.path.join(depth_tiff_folder, tiff_prefix_map[prefix])

        pcd_dst = os.path.join(output_pcd_folder, prefix + '.pcd')
        jpg_dst = os.path.join(output_meta_folder, jpg_prefix_map[prefix])
        json_dst = os.path.join(output_meta_folder, json_prefix_map[prefix])
        tiff_dst = os.path.join(output_tiff_folder, tiff_prefix_map[prefix])

        shutil.copy2(pcd_src, pcd_dst)
        shutil.copy2(jpg_src, jpg_dst)
        shutil.copy2(json_src, json_dst)
        shutil.copy2(tiff_src, tiff_dst)

    print(f"复制完成，匹配了 {len(matched_prefixes)} 组文件。")

if __name__ == "__main__":
    image_json_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\0525保存xyz信息的jpg+0516修订后的json"
    pcd_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\pcd"
    depth_tiff_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\depth_tiff"

    output_pcd_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\匹配后的pcd"
    output_meta_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\匹配后的json_jpg"
    output_tiff_folder = r"\\192.168.50.222\share\zyh\5月彩色钢板数据汇总\513-515数据汇总\匹配后的depth_tiff"

    match_pcd_meta_tiff(image_json_folder, pcd_folder, depth_tiff_folder,
                        output_pcd_folder, output_meta_folder, output_tiff_folder)