import bpy
import os
import sys

def clean_scene():
    bpy.ops.object.select_all(action='SELECT')
    bpy.ops.object.delete()
    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)

def remove_empty_vertex_groups(obj, threshold=0.001):
    """Removes vertex groups with no weights or weights below threshold."""
    if obj.type != 'MESH':
        return

    # Ensure we are in object mode
    bpy.context.view_layer.objects.active = obj

    # Dictionary to track max weight per group
    group_max_weights = {g.index: 0.0 for g in obj.vertex_groups}

    # Iterate over vertices to find actual max weights
    for v in obj.data.vertices:
        for g in v.groups:
            if g.group in group_max_weights:
                group_max_weights[g.group] = max(group_max_weights[g.group], g.weight)

    # Remove groups that don't meet the threshold
    groups_to_remove = [obj.vertex_groups[idx] for idx, max_w in group_max_weights.items() if max_w < threshold]

    for g in groups_to_remove:
        obj.vertex_groups.remove(g)

    print(f"Cleaned {len(groups_to_remove)} empty/low-weight groups from {obj.name}")

def process_batch():
    try:
        args = sys.argv[sys.argv.index("--") + 1:]
        clothing_path, lib_directory, out_dir = args[0], args[1], args[2]
    except (IndexError, ValueError):
        print("Usage: blender -b -P script.py -- <source_blend> <lib_dir> <out_dir>")
        return

    if not os.path.exists(out_dir): os.makedirs(out_dir)

    # 1. Identify all clothing in the source file
    with bpy.data.libraries.load(clothing_path) as (data_from, data_to):
        all_clothing_names = data_from.objects

    # Start with a fresh scene
    clean_scene()

    for name in all_clothing_names:
        # 2. Append the clothing item
        with bpy.data.libraries.load(clothing_path) as (data_from, data_to):
            data_to.objects = [name]
        
        for obj in data_to.objects:
            if obj: bpy.context.collection.objects.link(obj)

        clothing = bpy.data.objects.get(name)
        if not clothing or clothing.type != 'MESH': continue

        # Get properties
        sex = clothing.get("ref_sex")
        age = clothing.get("ref_age")
        ref_mesh_name = clothing.get("ref_clothing")

        if not all([sex, age, ref_mesh_name]):
            print(f"Skipping {name}: Missing properties")
            continue

        # 3. Locate Reference Library
        target_lib_name = f"normal_{age}_{sex}.blend"
        target_lib_path = os.path.join(lib_directory, target_lib_name)
        rig_name = str(sex)

        if not os.path.exists(target_lib_path):
            if target_lib_name == "normal_adult_male.blend":
                target_lib_name = "edited-normal-male.blend"
            elif target_lib_name == "normal_adult_female.blend":
                target_lib_name = "edited-normal-female.blend"
            target_lib_path = os.path.join(lib_directory, target_lib_name)
            if not os.path.exists(target_lib_path):
                print(f"Error: Library {target_lib_path} not found")
                continue

        # 4. Append Weights Source and Rig
        with bpy.data.libraries.load(target_lib_path) as (data_from, data_to):
            data_to.objects = [ref_mesh_name, rig_name]
        
        for obj in data_to.objects:
            if obj: bpy.context.collection.objects.link(obj)

        source_mesh = bpy.data.objects.get(ref_mesh_name)
        rig = bpy.data.objects.get(rig_name)

        # ---- Detect hair with its own skeleton ----
        has_own_arm = False
        own_arm_name = None
        own_arm = None
        print(f"  Checking {name} modifiers for own armature:")
        for mod in clothing.modifiers:
            print(f"    type={mod.type} object={mod.object}")
            if mod.type == 'ARMATURE':
                arm = mod.object
                if arm is None:
                    print(f"    mod.object is None, trying by name...")
                    continue
                print(f"    arm.name='{arm.name}' rig_name='{rig_name}'")
                if arm.name != rig_name:
                    has_own_arm = True
                    own_arm_name = arm.name
                    own_arm = arm
                    break
        if has_own_arm and own_arm:
            clothing["has_own_armature"] = True
            clothing["own_armature_name"] = own_arm_name
            # Ensure hair armature is linked to scene so parent/modifier
            # references survive the save.
            try:
                bpy.context.collection.objects.link(own_arm)
            except RuntimeError:
                pass  # Already linked
            print(f"  DETECTED own armature '{own_arm_name}' on {name}")
        else:
            print(f"  No own armature on {name} (rig='{rig_name}')")
        # -------------------------------------------

        # 5. Prep Objects (Apply Scale & Clear Animation)
        bpy.context.view_layer.objects.active = clothing
        bpy.ops.object.transform_apply(location=False, rotation=True, scale=True)
        
        # Don't clear animation on the hair armature if it has
        # its own skeleton (we want to preserve hair animations)
        items_to_clear = [clothing]
        if not has_own_arm:
            items_to_clear.append(rig)
        for item in items_to_clear:
            if item.animation_data: item.animation_data_clear()
            if item.type == 'ARMATURE':
                bpy.context.view_layer.objects.active = item
                bpy.ops.object.mode_set(mode='POSE')
                bpy.ops.pose.transforms_clear()
                bpy.ops.object.mode_set(mode='OBJECT')

        # Also clear the body rig's animation (always)
        if has_own_arm:
            for item in [rig]:
                if item.animation_data: item.animation_data_clear()
                if item.type == 'ARMATURE':
                    bpy.context.view_layer.objects.active = item
                    bpy.ops.object.mode_set(mode='POSE')
                    bpy.ops.pose.transforms_clear()
                    bpy.ops.object.mode_set(mode='OBJECT')

        # Set garment_id from object name if not already set
        if "garment_id" not in clothing:
            clothing["garment_id"] = clothing.name

        # 6. Weight Transfer
        if has_own_arm:
            # Hair with its own skeleton already has correct vertex
            # groups for its hair bones; body weight transfer would
            # destroy them.
            print(f"Skipped weight transfer for {name}: has own skeleton")
        else:
            clothing.vertex_groups.clear()
            dt = clothing.modifiers.new(name="WT", type='DATA_TRANSFER')
            dt.object = source_mesh
            dt.use_vert_data = True
            dt.data_types_verts = {'VGROUP_WEIGHTS'}
            dt.layers_vgroup_select_src = 'ALL'
            dt.vert_mapping = 'POLYINTERP_NEAREST'
            
            bpy.context.view_layer.objects.active = clothing
            # "Generate Data Layers" step
            bpy.ops.object.datalayout_transfer(modifier=dt.name)
            bpy.ops.object.modifier_apply(modifier=dt.name)

            remove_empty_vertex_groups(clothing, threshold=0.001)

        # 7. Final Parenting
        if has_own_arm:
            # Hair with its own skeleton: keep the existing Armature
            # modifier (which points to the hair skeleton).  Just
            # parent to the hair armature; do NOT add a body
            # Armature modifier.
            clothing.parent = own_arm
            print(f"Parented {name} to hair armature '{own_arm_name}'")
        else:
            clothing.parent = rig
            arm_mod = clothing.modifiers.new(name="Armature", type='ARMATURE')
            arm_mod.object = rig

        # 8. Cleanup Reference Mesh (keep the Rig!)
        bpy.data.objects.remove(source_mesh, do_unlink=True)
        print(f"Processed: {name}")

    # 9. Save as single file named after source + _weighted
    source_filename = os.path.splitext(os.path.basename(clothing_path))[0]
    final_save_path = os.path.join(out_dir, f"{source_filename}_weighted.blend")
    
    # Purge any remaining junk before final save
    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)
    bpy.ops.wm.save_as_mainfile(filepath=final_save_path)
    print(f"\n--- ALL DONE ---")
    print(f"Saved to: {final_save_path}")

if __name__ == "__main__":
    process_batch()
