import bpy
import sys
import os

def propagate_shape_keys(blend_path, output_path):
    """Copy missing shape keys from Body_shapes to all body part objects."""
    print(f"Propagating shape keys from {blend_path} to {output_path}")
    
    bpy.ops.wm.open_mainfile(filepath=blend_path)
    
    # Find source object with all shape keys
    source_obj = None
    for name in ['Body_shapes', 'Body']:
        if name in bpy.data.objects and bpy.data.objects[name].type == 'MESH':
            obj = bpy.data.objects[name]
            if obj.data.shape_keys and len(obj.data.shape_keys.key_blocks) > 1:
                source_obj = obj
                print(f"Found source object: {name}")
                break
    
    if not source_obj:
        print("ERROR: No source object with shape keys found")
        return False
    
    source_keys = source_obj.data.shape_keys.key_blocks
    source_basis = source_keys['Basis'].data
    
    modified = False
    for obj in bpy.data.objects:
        if obj.type != 'MESH':
            continue
        # Only process body part objects
        if not all(p in obj for p in ['age', 'sex', 'slot']):
            continue
        # Skip the source object itself
        if obj == source_obj:
            continue
        
        if not obj.data.shape_keys:
            # Create basis if missing
            obj.shape_key_add(name='Basis', from_mix=False)
        
        existing_names = {kb.name for kb in obj.data.shape_keys.key_blocks}
        
        for sk in source_keys:
            if sk.name == 'Basis':
                continue
            if sk.name in existing_names:
                print(f"  {obj.name}: already has '{sk.name}', skipping")
                continue
            
            # Copy shape key
            print(f"  {obj.name}: adding '{sk.name}' from {source_obj.name}")
            new_sk = obj.shape_key_add(name=sk.name, from_mix=False)
            
            # Transfer vertex positions relative to basis
            # We need to map source vertices to target vertices
            # For now, assume same topology (body parts were separated from same mesh)
            num_verts = min(len(new_sk.data), len(sk.data), len(source_basis))
            for i in range(num_verts):
                # The shape key position in source
                src_pos = sk.data[i].co
                src_basis_pos = source_basis[i].co
                # Offset from basis
                offset = src_pos - src_basis_pos
                # Apply to target using target's basis position
                new_sk.data[i].co = obj.data.shape_keys.key_blocks['Basis'].data[i].co + offset
            
            new_sk.data.update()
            modified = True
        
        if modified:
            obj.data.update_tag()
    
    bpy.context.view_layer.update()
    bpy.ops.wm.save_as_mainfile(filepath=output_path)
    print(f"Saved: {output_path}")
    return True

if __name__ == "__main__":
    try:
        args = sys.argv[sys.argv.index("--") + 1:]
        if len(args) >= 2:
            propagate_shape_keys(args[0], args[1])
        else:
            print("Usage: blender -b -P script.py -- <input.blend> <output.blend>")
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)
