import bpy

def separate_geometry_objects():
    # Get the active armature
    active_armature = bpy.context.active_object
    if not active_armature or active_armature.type != 'ARMATURE':
        print("Error: No active armature selected")
        return
    
    print(f"Working with armature: {active_armature.name}")
    
    # Find the collection that contains the armature
    armature_collection = None
    for collection in bpy.data.collections:
        if active_armature.name in collection.objects:
            armature_collection = collection
            break
    
    # If armature is in scene collection, use scene name
    if not armature_collection:
        if active_armature.name in bpy.context.scene.collection.objects:
            armature_collection_name = bpy.context.scene.name
        else:
            armature_collection_name = "Scene"
    else:
        armature_collection_name = armature_collection.name
    
    print(f"Armature is in collection: {armature_collection_name}")
    
    # Create new collection name
    geo_collection_name = f"GEO-{armature_collection_name}"
    
    # Check if the collection already exists
    geo_collection = bpy.data.collections.get(geo_collection_name)
    if not geo_collection:
        # Create new collection
        geo_collection = bpy.data.collections.new(geo_collection_name)
        
        # Link to the same collection as the armature
        if armature_collection:
            armature_collection.children.link(geo_collection)
            print(f"Created new collection: {geo_collection_name} inside {armature_collection.name}")
        else:
            bpy.context.scene.collection.children.link(geo_collection)
            print(f"Created new collection: {geo_collection_name} under scene collection")
    else:
        print(f"Using existing collection: {geo_collection_name}")
    
    # Set collection color to orange
    geo_collection.color_tag = 'COLOR_02'  # Orange color tag
    print(f"Set {geo_collection_name} color to orange")
    
    # Create subcollections inside GEO collection
    accessories_collection_name = "Accessories"
    clothing_collection_name = "Clothing"
    body_collection_name = "Body"
    
    # Create Accessories subcollection
    accessories_collection = bpy.data.collections.get(accessories_collection_name)
    if not accessories_collection:
        accessories_collection = bpy.data.collections.new(accessories_collection_name)
    
    # Ensure it's linked to GEO collection
    if accessories_collection_name not in [child.name for child in geo_collection.children]:
        geo_collection.children.link(accessories_collection)
        print(f"Created Accessories subcollection inside {geo_collection_name}")
    
    # Create Clothing subcollection
    clothing_collection = bpy.data.collections.get(clothing_collection_name)
    if not clothing_collection:
        clothing_collection = bpy.data.collections.new(clothing_collection_name)
    
    if clothing_collection_name not in [child.name for child in geo_collection.children]:
        geo_collection.children.link(clothing_collection)
        print(f"Created Clothing subcollection inside {geo_collection_name}")
    
    # Create Body subcollection
    body_collection = bpy.data.collections.get(body_collection_name)
    if not body_collection:
        body_collection = bpy.data.collections.new(body_collection_name)
    
    if body_collection_name not in [child.name for child in geo_collection.children]:
        geo_collection.children.link(body_collection)
        print(f"Created Body subcollection inside {geo_collection_name}")
    
    # Find all objects parented to the armature
    parented_objects = []
    for obj in bpy.data.objects:
        if obj.parent == active_armature:
            parented_objects.append(obj)
    
    print(f"Found {len(parented_objects)} objects parented to armature")
    
    # Organize objects by category
    body_moved_count = 0
    accessories_moved_count = 0
    fallback_moved_count = 0
    
    for obj in parented_objects:
        # Remove object from all collections it's currently in
        for collection in obj.users_collection:
            collection.objects.unlink(obj)
        
        # Check if object has CC_ prefix (goes to Body)
        if obj.name.startswith('CC_'):
            body_collection.objects.link(obj)
            body_moved_count += 1
            print(f"Moved {obj.name} to Body")
        # Check if object should go to Accessories by name (only specific items)
        elif obj.name in ['Device', 'device-band', 'Finger-Scanner', 'Lanyard', 'Vest']:
            accessories_collection.objects.link(obj)
            accessories_moved_count += 1
            print(f"Moved {obj.name} to Accessories")
        # Everything else goes to main GEO collection
        else:
            geo_collection.objects.link(obj)
            fallback_moved_count += 1
            print(f"Moved {obj.name} to {geo_collection_name}")
    
    # Also find any standalone CC_ objects that aren't parented to armature
    standalone_cc_objects = []
    for obj in bpy.data.objects:
        if obj.name.startswith('CC_') and obj.parent != active_armature:
            standalone_cc_objects.append(obj)
    
    if standalone_cc_objects:
        print(f"Found {len(standalone_cc_objects)} standalone CC_ objects")
        for obj in standalone_cc_objects:
            # Remove object from all collections it's currently in
            for collection in obj.users_collection:
                collection.objects.unlink(obj)
            
            # Add to Body collection
            body_collection.objects.link(obj)
            body_moved_count += 1
            print(f"Moved standalone {obj.name} to Body")
    
    print(f"Successfully organized objects:")
    print(f"  - {accessories_moved_count} objects moved to Accessories")
    print(f"  - {body_moved_count} objects moved to Body")
    print(f"  - {fallback_moved_count} objects moved to {geo_collection_name} (fallback)")
    print(f"  - Clothing subcollection created (objects to be moved manually)")

# Execute the operation
separate_geometry_objects()