import bpy

def add_body_masks():
    # Find the CC_Base_Body object
    body_obj = bpy.data.objects.get('CC_Base_Body')
    if not body_obj:
        print("Error: CC_Base_Body object not found")
        return
    
    print(f"Found body object: {body_obj.name}")
    
    # Define vertex groups for each mask
    main_mask_groups = [
        # Arms
        'DEF-shoulder.L', 'DEF-shoulder.R',
        'DEF-upper_arm.L', 'DEF-upper_arm.L.001', 'DEF-upper_arm.R', 'DEF-upper_arm.R.001',
        'DEF-elbow_share.L', 'DEF-elbow_share.R',
        'DEF-forearm.L', 'DEF-forearm.L.001', 'DEF-forearm.R', 'DEF-forearm.R.001',
        # Head and chest (upper torso)
        'DEF-spine.003', 'DEF-spine.004', 'DEF-spine.005', 'DEF-spine.006',
        'DEF-jaw', 'DEF-jaw.L', 'DEF-jaw.R', 'DEF-jaw.L.001', 'DEF-jaw.R.001',
        'DEF-forehead.L', 'DEF-forehead.R', 'DEF-forehead.L.001', 'DEF-forehead.R.001', 'DEF-forehead.L.002', 'DEF-forehead.R.002',
        'DEF-temple.L', 'DEF-temple.R',
        'DEF-brow.B.L', 'DEF-brow.B.L.001', 'DEF-brow.B.L.002', 'DEF-brow.B.L.003',
        'DEF-brow.B.R', 'DEF-brow.B.R.001', 'DEF-brow.B.R.002', 'DEF-brow.B.R.003',
        'DEF-brow.T.L', 'DEF-brow.T.L.001', 'DEF-brow.T.L.002', 'DEF-brow.T.L.003',
        'DEF-brow.T.R', 'DEF-brow.T.R.001', 'DEF-brow.T.R.002', 'DEF-brow.T.R.003',
        'DEF-lid.B.L', 'DEF-lid.B.L.001', 'DEF-lid.B.L.002', 'DEF-lid.B.L.003',
        'DEF-lid.B.R', 'DEF-lid.B.R.001', 'DEF-lid.B.R.002', 'DEF-lid.B.R.003',
        'DEF-lid.T.L', 'DEF-lid.T.L.001', 'DEF-lid.T.L.002', 'DEF-lid.T.L.003',
        'DEF-lid.T.R', 'DEF-lid.T.R.001', 'DEF-lid.T.R.002', 'DEF-lid.T.R.003',
        'DEF-ear.L', 'DEF-ear.L.001', 'DEF-ear.L.002', 'DEF-ear.L.003', 'DEF-ear.L.004',
        'DEF-ear.R', 'DEF-ear.R.001', 'DEF-ear.R.002', 'DEF-ear.R.003', 'DEF-ear.R.004',
        'DEF-chin', 'DEF-chin.001', 'DEF-chin.L', 'DEF-chin.R',
        'DEF-cheek.T.L', 'DEF-cheek.T.R', 'DEF-cheek.T.L.001', 'DEF-cheek.T.R.001',
        'DEF-cheek.B.L', 'DEF-cheek.B.R', 'DEF-cheek.B.L.001', 'DEF-cheek.B.R.001',
        'DEF-nose', 'DEF-nose.L', 'DEF-nose.R', 'DEF-nose.001', 'DEF-nose.002', 'DEF-nose.003', 'DEF-nose.004',
        'DEF-nose.L.001', 'DEF-nose.R.001',
        'DEF-lip.B.L', 'DEF-lip.B.R', 'DEF-lip.B.L.001', 'DEF-lip.B.R.001',
        'DEF-lip.T.L', 'DEF-lip.T.R', 'DEF-lip.T.L.001', 'DEF-lip.T.R.001',
        'DEF-tongue', 'DEF-tongue.001', 'DEF-tongue.002',
        'DEF-breast.L', 'DEF-breast.R', 'DEF-breast_twist.L', 'DEF-breast_twist.R'
    ]
    
    hand_groups = [
        # Hands
        'DEF-hand.L', 'DEF-hand.R',
        'DEF-f_pinky.01.L', 'DEF-f_pinky.02.L', 'DEF-f_pinky.03.L',
        'DEF-f_ring.01.L', 'DEF-f_ring.02.L', 'DEF-f_ring.03.L',
        'DEF-f_middle.01.L', 'DEF-f_middle.02.L', 'DEF-f_middle.03.L',
        'DEF-f_index.01.L', 'DEF-f_index.02.L', 'DEF-f_index.03.L',
        'DEF-thumb.01.L', 'DEF-thumb.02.L', 'DEF-thumb.03.L',
        'DEF-f_pinky.01.R', 'DEF-f_pinky.02.R', 'DEF-f_pinky.03.R',
        'DEF-f_ring.01.R', 'DEF-f_ring.02.R', 'DEF-f_ring.03.R',
        'DEF-f_middle.01.R', 'DEF-f_middle.02.R', 'DEF-f_middle.03.R',
        'DEF-f_index.01.R', 'DEF-f_index.02.R', 'DEF-f_index.03.R',
        'DEF-thumb.01.R', 'DEF-thumb.02.R', 'DEF-thumb.03.R'
    ]
    
    # Create vertex groups for masks
    def create_mask_vertex_group(obj, group_name, vertex_group_names):
        # Remove existing group if it exists
        existing_group = obj.vertex_groups.get(group_name)
        if existing_group:
            obj.vertex_groups.remove(existing_group)
        
        # Create new vertex group
        mask_group = obj.vertex_groups.new(name=group_name)
        
        # Get indices of source vertex groups
        source_group_indices = []
        for vg_name in vertex_group_names:
            vg = obj.vertex_groups.get(vg_name)
            if vg:
                source_group_indices.append(vg.index)
        
        if not source_group_indices:
            print(f"Warning: No source vertex groups found for {group_name}")
            return None
        
        # Add vertices that have weights in any of the source groups
        vertices_to_add = []
        for vert_idx, vert in enumerate(obj.data.vertices):
            for group in vert.groups:
                if group.group in source_group_indices and group.weight > 0.0:
                    vertices_to_add.append(vert_idx)
                    break
        
        if vertices_to_add:
            mask_group.add(vertices_to_add, 1.0, 'REPLACE')
            print(f"Created {group_name} vertex group with {len(vertices_to_add)} vertices")
        
        return mask_group
    
    # Create Main mask vertex group (head, arms, chest)
    main_mask_vg = create_mask_vertex_group(body_obj, "Main_Mask", main_mask_groups)
    
    # Create Hand mask vertex group (main + hands)
    hand_mask_vg = create_mask_vertex_group(body_obj, "Hand_Mask", main_mask_groups + hand_groups)
    
    # Add mask modifiers
    # Remove existing mask modifiers if they exist
    modifiers_to_remove = []
    for modifier in body_obj.modifiers:
        if modifier.type == 'MASK' and modifier.name in ['Main_Mask', 'Hand_Mask']:
            modifiers_to_remove.append(modifier)
    
    for modifier in modifiers_to_remove:
        body_obj.modifiers.remove(modifier)
        print(f"Removed existing {modifier.name} modifier")
    
    # Add Main mask modifier
    if main_mask_vg:
        main_mask_modifier = body_obj.modifiers.new(name="Main_Mask", type='MASK')
        main_mask_modifier.vertex_group = "Main_Mask"
        print("Added Main_Mask modifier (head, arms, chest)")
    
    # Add Hand mask modifier
    if hand_mask_vg:
        hand_mask_modifier = body_obj.modifiers.new(name="Hand_Mask", type='MASK')
        hand_mask_modifier.vertex_group = "Hand_Mask"
        print("Added Hand_Mask modifier (head, arms, chest + hands)")
    
    print("\nBody masking completed!")
    print("Main_Mask: Shows head, arms, and chest")
    print("Hand_Mask: Shows head, arms, chest, and hands")

# Execute the operation
add_body_masks()