bl_info = {
    "name": "Align Advanced",
    "author": "Antigravity",
    "version": (1, 0),
    "blender": (3, 0, 0),
    "location": "View3D > Sidebar > Edit Tab",
    "description": "Aligns selected objects to active object with advanced options",
    "warning": "",
    "wiki_url": "",
    "category": "Object",
}

import bpy
from mathutils import Matrix, Vector

class ALIGN_OT_advanced(bpy.types.Operator):
    """Align selected objects to the active object"""
    bl_idname = "object.align_advanced"
    bl_label = "Align Advanced"
    bl_options = {'REGISTER', 'UNDO'}

    align_position: bpy.props.BoolProperty(
        name="Align Position",
        description="Align position to active object",
        default=True
    )
    
    align_rotation: bpy.props.BoolProperty(
        name="Align Rotation",
        description="Align rotation to active object",
        default=True
    )
    
    affect_pivot: bpy.props.BoolProperty(
        name="Affect Pivot Only",
        description="Only move the origin, keeping geometry in place",
        default=False
    )

    @classmethod
    def poll(cls, context):
        return context.active_object is not None and len(context.selected_objects) >= 2

    def invoke(self, context, event):
        return context.window_manager.invoke_props_dialog(self)

    def execute(self, context):
        # Update the view layer to ensure all matrices are fresh
        context.view_layer.update()

        active_obj = context.active_object
        selected_objs = [o for o in context.selected_objects if o != active_obj]
        
        if not active_obj:
            self.report({'ERROR'}, "No active object selected")
            return {'CANCELLED'}

        # Get active object's world transforms
        target_matrix = active_obj.matrix_world
        target_loc, target_rot, _ = target_matrix.decompose()
        
        # Keep track of modified data to prevent double-transforming shared meshes
        processed_data = set()
        
        for obj in selected_objs:
            # Get current transforms
            orig_matrix = obj.matrix_world.copy()
            orig_loc, orig_rot, orig_scale = orig_matrix.decompose()
            
            # Determine new transforms based on settings
            new_loc = target_loc if self.align_position else orig_loc
            new_rot = target_rot if self.align_rotation else orig_rot
            
            # Construct new world matrix
            # Using original scale
            mat_loc = Matrix.Translation(new_loc)
            mat_rot = new_rot.to_matrix().to_4x4()
            mat_scale = Matrix.Diagonal(orig_scale).to_4x4()
            
            new_matrix = mat_loc @ mat_rot @ mat_scale
            
            if self.affect_pivot:
                # If affecting pivot only, we must ensure vertices and children keep world position
                
                # Check for shared data (Linked Duplicates)
                if obj.data in processed_data:
                    # If we already moved this mesh for another object, we cannot move it again 
                    # significantly without breaking the previous object's visual alignment.
                    # We just update the pivot for this object.
                    obj.matrix_world = new_matrix
                    continue

                # 1. Transform Mesh Data
                # We want: new_matrix @ new_local = orig_matrix @ old_local
                # So: new_local = new_matrix.inv @ orig_matrix @ old_local
                if obj.data and hasattr(obj.data, "transform"):
                    try:
                        transform_mat = new_matrix.inverted() @ orig_matrix
                        obj.data.transform(transform_mat)
                        obj.data.update()
                        processed_data.add(obj.data)
                    except Exception as e:
                        print(f"Could not transform data for {obj.name}: {e}")

                # 2. Handle Children (keep them in place)
                # Store their current world matrices
                child_globals = {child: child.matrix_world.copy() for child in obj.children}
                
                # Apply the new matrix to the parent (the pivot move)
                obj.matrix_world = new_matrix
                
                # Restore children world positions
                for child, mat in child_globals.items():
                    child.matrix_world = mat
                    
            else:
                # Simple alignment (moves geometry with the object)
                obj.matrix_world = new_matrix

        return {'FINISHED'}

def menu_func(self, context):
    self.layout.separator()
    self.layout.operator("object.align_advanced", text="Align Advanced")

addon_keymaps = []

def register():
    bpy.utils.register_class(ALIGN_OT_advanced)
    bpy.types.VIEW3D_MT_object_context_menu.append(menu_func)
    
    # Add Shortcut CTRL + ALT + A
    wm = bpy.context.window_manager
    kc = wm.keyconfigs.addon
    if kc:
        km = kc.keymaps.new(name='Object Mode', space_type='EMPTY')
        kmi = km.keymap_items.new("object.align_advanced", 'A', 'PRESS', ctrl=True, alt=True)
        addon_keymaps.append((km, kmi))

def unregister():
    bpy.types.VIEW3D_MT_object_context_menu.remove(menu_func)
    bpy.utils.unregister_class(ALIGN_OT_advanced)
    
    # Remove Shortcut
    for km, kmi in addon_keymaps:
        km.keymap_items.remove(kmi)
    addon_keymaps.clear()

if __name__ == "__main__":
    register()