import os
import re
import shutil
import json  # 修复：补充json导入（原代码虽未用，但避免后续扩展报错）
import pandas as pd  # 修复：补充pandas导入（同理，确保依赖完整）
from collections import defaultdict

def extract_merge_key(filename):
    """提取合并键：开头连续数字为键（用于分组），字母开头返回原文件名并标记"""
    name_without_ext = os.path.splitext(filename)[0]
    # 匹配开头连续数字（如"123_bolt"→"123"，"A001"无匹配）
    num_match = re.match(r'^\d+', name_without_ext)
    if num_match:
        return num_match.group(), 'num_start'  # 数字开头：返回数字键+类型
    else:
        return name_without_ext, 'letter_start'  # 字母开头：返回原名+类型

def merge_obj_by_prefix(input_dir, output_dir):
    """
    按.obj前缀合并：每组复制排序后第一个文件到输出目录，字母开头文件单独打印
    """
    # 1. 初始化存储：数字前缀→文件路径列表，字母开头文件列表
    num_groups = defaultdict(list)  # key: 数字前缀, value: [文件路径1, 文件路径2, ...]
    letter_files = []               # 存储字母开头的文件名（仅记录，不复制）

    # 2. 扫描输入目录所有.obj文件（递归包含子目录）
    print(f"开始扫描 {input_dir} 下的.obj文件...\n")
    for root, _, files in os.walk(input_dir):
        for file in files:
            if file.lower().endswith('.obj'):  # 兼容 .OBJ 大写后缀
                file_path = os.path.join(root, file)
                merge_key, start_type = extract_merge_key(file)
                
                if start_type == 'num_start':
                    # 数字开头：加入对应分组
                    num_groups[merge_key].append(file_path)
                    print(f"  数字开头：{file} → 归入组【{merge_key}】")
                else:
                    # 字母开头：加入未合并列表
                    letter_files.append(file)
                    print(f"  字母开头：{file} → 不合并（仅记录）")

    # 3. 检查是否有可处理的文件
    if not num_groups and not letter_files:
        print(f"警告：在 {input_dir} 中未找到任何.obj文件！")
        return

    # 4. 创建输出目录（清空原有内容，避免旧文件干扰）
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)  # 删除原有目录及文件
    os.makedirs(output_dir, exist_ok=True)
    print(f"\n已创建干净的输出目录：{output_dir}")

    # 5. 处理每个数字分组：排序后复制第一个文件到输出目录
    copied_count = 0
    print(f"\n开始处理合并组（共 {len(num_groups)} 组）：")
    for group_key in sorted(num_groups.keys()):  # 按数字前缀排序（如123→456）
        file_paths = num_groups[group_key]
        # 对组内文件按文件名排序（确保每次复制顺序一致）
        sorted_files = sorted(file_paths, key=lambda x: os.path.basename(x))
        first_file = sorted_files[0]  # 取排序后第一个文件
        first_filename = os.path.basename(first_file)
        
        # 复制文件到输出目录（保留原文件名）
        dest_path = os.path.join(output_dir, first_filename)
        shutil.copy2(first_file, dest_path)  # copy2 保留文件元信息（修改时间等）
        copied_count += 1
        
        print(f"  组【{group_key}】：复制排序后第一个文件 → {first_filename}")

    # 6. 单独打印字母开头的未合并文件
    print(f"\n=== 字母开头的未合并零件列表（共 {len(letter_files)} 个）===")
    if letter_files:
        for idx, file in enumerate(sorted(letter_files), 1):  # 排序后打印，更整齐
            print(f"  {idx}. {file}")
    else:
        print("  无字母开头的未合并零件")

    # 7. 输出最终统计
    print(f"\n=== 处理完成统计 ===")
    print(f"  合并组总数：{len(num_groups)} 组")
    print(f"  成功复制文件数：{copied_count} 个（输出目录：{output_dir}）")
    print(f"  未合并文件数：{len(letter_files)} 个（仅记录，未复制）")

if __name__ == "__main__":
    # --------------------------
    # 配置参数（请根据实际情况修改）
    # --------------------------
    INPUT_DIR = "/public/home/wangshuo/gap/assembly/data/part_obj_300"  # .obj源文件目录（递归扫描）
    OUTPUT_DIR = "/public/home/wangshuo/gap/assembly/data/part_obj_300_merge_by_prefix"  # 合并后文件输出目录（仅存复制的文件）
    
    # 执行合并逻辑
    merge_obj_by_prefix(INPUT_DIR, OUTPUT_DIR)
    print(f"\n✅ .obj零件合并复制完成！输出目录仅保留每组排序后第一个文件。")