import os
import shutil
import time

def merged_obj_for_group(input_folder,output_folder):
    """"""
    group_size = 5
    obj_pid_list = os.listdir(input_folder)
    group_obj_list = [obj_pid_list[i:i + group_size] for i in range(0, len(obj_pid_list), group_size)]
    print(group_obj_list)
    for group_obj in group_obj_list:
        print(group_obj)
        group_pid = "_".join(group_obj)
        print(group_pid)

        output_group_folder = os.path.join(output_folder,group_pid)
        os.makedirs(output_group_folder, exist_ok=True)

        #input_root_folder = "/data/datasets_20t/obj_merger_test_data/"
        #output_folder = "/data/datasets_20t/obj_merger_result/"
        output_obj = os.path.join(output_group_folder, f"{group_pid}.obj")
        output_mtl = os.path.join(output_group_folder, f"{group_pid}.mtl")

        # 初始化
        merged_obj = []
        merged_mtl = []
        texture_files = set()
        material_offset = 0
        vertex_offset = 0
        texture_offset = 0
        normal_offset = 0
        current_materials = {}

        # 遍历每个子文件夹
        for folder_name in group_obj:
            folder_path = os.path.join(input_folder, folder_name)
            if not os.path.isdir(folder_path):
                continue

            obj_file = None
            mtl_file = None
            texture_file = None

            # 寻找 .obj、.mtl 和 .jpg 文件
            for file_name in os.listdir(folder_path):
                if file_name.endswith(".obj"):
                    obj_file = os.path.join(folder_path, file_name)
                elif file_name.endswith(".mtl"):
                    mtl_file = os.path.join(folder_path, file_name)
                elif file_name.endswith(".jpg"):
                    texture_file = os.path.join(folder_path, file_name)

            # 跳过不完整的文件夹
            if not obj_file or not mtl_file or not texture_file:
                print(f"跳过不完整的文件夹：{folder_path}")
                continue

            # 读取 .obj 文件
            with open(obj_file, "r") as obj_f:
                obj_lines = obj_f.readlines()


            for line in obj_lines:
                if line.startswith("mtllib"):
                    # 替换材质文件名

                    merged_obj.append(f"mtllib {os.path.basename(output_mtl)}\n")

                elif line.startswith("usemtl"):
                    # 重命名材质名称，避免冲突
                    original_material = line.split()[1]
                    new_material = f"{original_material}_{material_offset}"
                    print(f"original_material---{original_material}")
                    print(f"new_material---{new_material}")
                    merged_obj.append(f"usemtl {new_material}\n")
                    current_materials[original_material] = new_material
                elif line.startswith("v "):  # 顶点
                    vertex = line.split()[1:]
                    merged_obj.append(f"v {' '.join(vertex)}\n")
                elif line.startswith("vt "):  # 纹理坐标
                    texture = line.split()[1:]
                    merged_obj.append(f"vt {' '.join(texture)}\n")
                elif line.startswith("vn "):  # 法线
                    normal = line.split()[1:]
                    merged_obj.append(f"vn {' '.join(normal)}\n")
                elif line.startswith("f "):  # 面数据
                    face = line.split()[1:]
                    updated_face = []
                    for vertex in face:
                        indices = vertex.split("/")
                        indices = [
                            str(int(indices[0]) + vertex_offset) if indices[0] else "",
                            str(int(indices[1]) + texture_offset) if len(indices) > 1 and indices[1] else "",
                            str(int(indices[2]) + normal_offset) if len(indices) > 2 and indices[2] else "",
                        ]
                        updated_face.append("/".join(indices))
                    merged_obj.append(f"f {' '.join(updated_face)}\n")

            # 更新偏移量
            vertex_offset += sum(1 for line in obj_lines if line.startswith("v "))
            texture_offset += sum(1 for line in obj_lines if line.startswith("vt "))
            normal_offset += sum(1 for line in obj_lines if line.startswith("vn "))

            # 读取 .mtl 文件
            with open(mtl_file, "r") as mtl_f:
                mtl_lines = mtl_f.readlines()

            for line in mtl_lines:
                if line.startswith("newmtl"):
                    # 重命名材质
                    original_material = line.split()[1]
                    new_material = current_materials.get(original_material, original_material)
                    merged_mtl.append(f"newmtl {new_material}\n")
                elif line.startswith(("map_Kd", "map_Ka", "map_bump")):
                    # 替换贴图路径为相对路径
                    texture_name = os.path.basename(texture_file)
                    merged_mtl.append(f"{line.split()[0]} {texture_name}\n")
                    texture_files.add(texture_file)
                else:
                    merged_mtl.append(line)

            material_offset += 1

        # 写入合并后的 .obj 和 .mtl 文件
        with open(output_obj, "w") as obj_out:
            obj_out.writelines(merged_obj)

        with open(output_mtl, "w") as mtl_out:
            mtl_out.writelines(merged_mtl)
        print(f"texture_files====={texture_files}")
        # 将纹理文件复制到输出目录
        for texture_file in texture_files:
            shutil.copy(texture_file, output_group_folder)

        print(f"合并完成：{output_obj} 和 {output_mtl}")
        print(f"纹理文件已复制到：{output_group_folder}")

def move_compact_obj_to_file(input_folder,output_folder):
    """"""
    group_size = 50
    obj_pid_list = os.listdir(input_folder)
    group_obj_list = [obj_pid_list[i:i + group_size] for i in range(0, len(obj_pid_list), group_size)]
    print(group_obj_list)
    for group_obj in group_obj_list:
        print(f"group_obj{group_obj}")
        out_obj_file_name = group_obj[0]+"_"+group_obj[-1]
        print(f"out_obj_file_name:::{out_obj_file_name}")
        group_out_put_dir = os.path.join(output_folder,out_obj_file_name)
        os.makedirs(group_out_put_dir,exist_ok=True)
        for obj_name in group_obj:
            original_obj_dir = os.path.join(input_folder,obj_name)
            for file_name in os.listdir(original_obj_dir):
                original_path = os.path.join(original_obj_dir,file_name)
                dis_path = os.path.join(group_out_put_dir,file_name)
                shutil.copy(original_path,dis_path)
    print("分组完成。")

if __name__ == '__main__':
    input_folder = "/data/datasets_20t/type_setting_test_data/print_compact_obj/"
    output_folder = "/data/datasets_20t/type_setting_test_data/obj_merger_result/"
    #merged_obj_for_group(input_folder,output_folder)
    move_compact_obj_to_file(input_folder, output_folder)