feat: undo for when an element gets deleted from media item deletion

2 files changed, 35 insertions, 14 deletions

diff --git a/apps/web/src/stores/media-store.ts b/apps/web/src/stores/media-store.ts
index 45f6ee04..650c0540 100644
--- a/apps/web/src/stores/media-store.ts
+++ b/apps/web/src/stores/media-store.ts
@@ -207,17 +207,29 @@ export const useMediaStore = create<MediaStore>((set, get) => ({
       removeElementFromTrack,
       removeElementFromTrackWithRipple,
       rippleEditingEnabled,
+      pushHistory,
     } = timeline;
 
-    // Iterate over a snapshot of tracks and their elements
+    // Find all elements that reference this media
+    const elementsToRemove: Array<{ trackId: string; elementId: string }> = [];
     for (const track of tracks) {
       for (const el of track.elements) {
         if (el.type === "media" && el.mediaId === id) {
-          if (rippleEditingEnabled) {
-            removeElementFromTrackWithRipple(track.id, el.id);
-          } else {
-            removeElementFromTrack(track.id, el.id);
-          }
+          elementsToRemove.push({ trackId: track.id, elementId: el.id });
+        }
+      }
+    }
+
+    // If there are elements to remove, push history once before batch removal
+    if (elementsToRemove.length > 0) {
+      pushHistory();
+
+      // Remove all elements without pushing additional history entries
+      for (const { trackId, elementId } of elementsToRemove) {
+        if (rippleEditingEnabled) {
+          removeElementFromTrackWithRipple(trackId, elementId, false);
+        } else {
+          removeElementFromTrack(trackId, elementId, false);
         }
       }
     }
diff --git a/apps/web/src/stores/timeline-store.ts b/apps/web/src/stores/timeline-store.ts
index 820025cc..6f9dfbd6 100644
--- a/apps/web/src/stores/timeline-store.ts
+++ b/apps/web/src/stores/timeline-store.ts
@@ -96,7 +96,11 @@ interface TimelineStore {
   removeTrack: (trackId: string) => void;
   removeTrackWithRipple: (trackId: string) => void;
   addElementToTrack: (trackId: string, element: CreateTimelineElement) => void;
-  removeElementFromTrack: (trackId: string, elementId: string) => void;
+  removeElementFromTrack: (
+    trackId: string,
+    elementId: string,
+    pushHistory?: boolean
+  ) => void;
   moveElementToTrack: (
     fromTrackId: string,
     toTrackId: string,
@@ -156,7 +160,8 @@ interface TimelineStore {
   ) => void;
   removeElementFromTrackWithRipple: (
     trackId: string,
-    elementId: string
+    elementId: string,
+    pushHistory?: boolean
   ) => void;
 
   // Computed values
@@ -553,13 +558,13 @@ export const useTimelineStore = create<TimelineStore>((set, get) => {
       get().selectElement(trackId, newElement.id);
     },
 
-    removeElementFromTrack: (trackId, elementId) => {
+    removeElementFromTrack: (trackId, elementId, pushHistory = true) => {
       const { rippleEditingEnabled } = get();
 
       if (rippleEditingEnabled) {
-        get().removeElementFromTrackWithRipple(trackId, elementId);
+        get().removeElementFromTrackWithRipple(trackId, elementId, pushHistory);
       } else {
-        get().pushHistory();
+        if (pushHistory) get().pushHistory();
         updateTracksAndSave(
           get()
             ._tracks.map((track) =>
@@ -577,12 +582,16 @@ export const useTimelineStore = create<TimelineStore>((set, get) => {
       }
     },
 
-    removeElementFromTrackWithRipple: (trackId, elementId) => {
+    removeElementFromTrackWithRipple: (
+      trackId,
+      elementId,
+      pushHistory = true
+    ) => {
       const { _tracks, rippleEditingEnabled } = get();
 
       if (!rippleEditingEnabled) {
         // If ripple editing is disabled, use regular removal
-        get().removeElementFromTrack(trackId, elementId);
+        get().removeElementFromTrack(trackId, elementId, pushHistory);
         return;
       }
 
@@ -591,7 +600,7 @@ export const useTimelineStore = create<TimelineStore>((set, get) => {
 
       if (!element || !track) return;
 
-      get().pushHistory();
+      if (pushHistory) get().pushHistory();
 
       const elementStartTime = element.startTime;
       const elementDuration =