import bpy
import sys
import os

# Import seam/normal fix functions from transfer_shape_keys.py
_script_dir = os.path.dirname(os.path.abspath(__file__))
if _script_dir not in sys.path:
    sys.path.insert(0, _script_dir)
from transfer_shape_keys import fix_seams_across_objects, fix_normals_across_objects

def process_append(source_files, output_path):
    required_props = {"age", "sex", "slot"}

    for file_path in source_files:
        if not os.path.exists(file_path):
            print(f"Warning: Source file not found: {file_path}")
            continue

        # ---- Pre-import materials and armatures from source blend ----
        with bpy.data.libraries.load(file_path, link=False) as (data_from, data_to):
            data_to.materials = data_from.materials
            data_to.objects = data_from.objects
        imported_materials = {}
        for mat in data_to.materials:
            if mat is None:
                continue
            existing = bpy.data.materials.get(mat.name)
            if existing and existing != mat:
                print(f"  Material '{mat.name}' already exists with "
                      f"different data; incoming material will "
                      f"replace it on appended objects.")
            imported_materials[mat.name] = mat

        # Also track imported armatures
        imported_armatures = {}
        for obj in data_to.objects:
            if obj is None:
                continue
            if obj.type == 'ARMATURE':
                imported_armatures[obj.name] = obj
        # ----------------------------------------------------

        for obj in data_to.objects:
            if obj is None: continue

            # Keep armatures from the source blend
            if obj.type == 'ARMATURE':
                bpy.context.collection.objects.link(obj)
                print(f"Imported armature '{obj.name}' from source blend")
                continue

            # Check criteria for body part meshes
            has_props = all(p in obj.keys() for p in required_props)
            if obj.type == 'MESH' and has_props:
                # 1. Link to the scene root temporarily
                bpy.context.collection.objects.link(obj)

                # ---- Remap material slots to imported versions ----
                for slot in obj.material_slots:
                    if slot.material and slot.material.name in imported_materials:
                        imported = imported_materials[slot.material.name]
                        if imported != slot.material:
                            print(f"  Remapping slot "
                                  f"'{slot.material.name}' -> "
                                  f"'{imported.name}' on '{obj.name}'")
                            slot.material = imported
                # ---------------------------------------------------

                # 2. Synchronize Names
                obj.data.name = obj.name

                # 3. Armature handling
                has_own_arm = obj.get("has_own_armature", False)
                if has_own_arm:
                    own_arm_name = obj.get("own_armature_name", "")
                    own_arm = bpy.data.objects.get(own_arm_name)
                    if not own_arm and own_arm_name in imported_armatures:
                        own_arm = imported_armatures[own_arm_name]
                    if own_arm:
                        # Keep hair's own armature; don't reparent
                        # to body armature.
                        # Fix the Armature modifier to point to the
                        # hair armature (combine_clothes may have
                        # changed it to the body rig).
                        arm_mod = next((m for m in obj.modifiers
                                        if m.type == 'ARMATURE'), None)
                        if arm_mod:
                            arm_mod.object = own_arm
                        obj.parent = own_arm
                        print(f"Processed {obj.name}: keeping own "
                              f"armature '{own_arm_name}'")
                        continue

                # Normal body part: parent to body armature
                arm_name = obj.get("sex")
                arm_obj = bpy.data.objects.get(arm_name)

                if arm_obj and arm_obj.type == 'ARMATURE':
                    # A. Handle Collections
                    for col in obj.users_collection:
                        col.objects.unlink(obj)
                    for col in arm_obj.users_collection:
                        col.objects.link(obj)

                    # B. Parent to Armature
                    obj.parent = arm_obj

                    # C. Handle Armature Modifier
                    arm_mod = next((m for m in obj.modifiers if m.type == 'ARMATURE'), None)
                    if not arm_mod:
                        arm_mod = obj.modifiers.new(name="Armature", type='ARMATURE')
                    arm_mod.object = arm_obj
                    print(f"Processed {obj.name}: Parented and Modset to {arm_name}")
                else:
                    print(f"Warning: Armature '{arm_name}' not found for {obj.name}")
            elif obj.type == 'MESH':
                missing = [p for p in required_props if p not in obj.keys()]
                print(f"WARNING: Mesh object '{obj.name}' is missing required properties {missing}.")
                print(f"  Treating as helper/reference object (not a body part).")
                print(f"  Available properties: {[k for k in obj.keys() if not k.startswith('_')]}")
            else:
                bpy.data.objects.remove(obj, do_unlink=True)

    # 4. Recursive Purge of all unlinked data (Materials, Textures, Meshes)
    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)

    # 5. Fix seams across ALL body parts in the consolidated scene.
    print("\nFixing seams across all consolidated body parts...")
    fix_seams_across_objects()
    fix_normals_across_objects()
    print("Consolidated seam/normal fix complete.\n")

    # Save
    bpy.ops.wm.save_as_mainfile(filepath=output_path)

if __name__ == "__main__":
    try:
        args = sys.argv[sys.argv.index("--") + 1:]
        if len(args) >= 2:
            *sources, output = args
            process_append(sources, output)
    except ValueError:
        print("Error: Use '--' to separate Blender args from script args.")
