From 91d35f5343d727bbea8d4cdb7796702240deda99 Mon Sep 17 00:00:00 2001
From: Tanish Manku <tanish2k09.dev@gmail.com>
Date: Wed, 27 Aug 2025 12:27:14 -0700
Subject: [PATCH] Move history stack state to zustand

---
 .../app/src/pages/PageEdit/HistoryStack.tsx   | 28 ++++++-------------
 .../app/src/pages/PageEdit/PageEdit.tsx       |  9 ++++--
 .../src/pages/PageEdit/stores/EditState.ts    | 19 ++++++++++++-
 3 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/frontend/apps/editor3d/app/src/pages/PageEdit/HistoryStack.tsx b/frontend/apps/editor3d/app/src/pages/PageEdit/HistoryStack.tsx
index f0bd918e88..7dbfdf5a55 100644
--- a/frontend/apps/editor3d/app/src/pages/PageEdit/HistoryStack.tsx
+++ b/frontend/apps/editor3d/app/src/pages/PageEdit/HistoryStack.tsx
@@ -23,21 +23,18 @@ export interface ImageBundle {
 interface HistoryStackProps {
   onClear: () => void;
   onImageSelect?: (image: BaseSelectorImage) => void;
-  startingBundles: ImageBundle[];
+  imageBundles: ImageBundle[];
+  onNewImageBundle?: (newBundle: ImageBundle) => void;
+  selectedImageToken?: string;
 }
 
 export const HistoryStack = ({
   onClear,
-  onImageSelect = () => {},
-  startingBundles = [],
+  onImageSelect = () => { },
+  imageBundles,
+  onNewImageBundle = () => { },
+  selectedImageToken
 }: HistoryStackProps) => {
-  const [imageBundles, setImageBundles] =
-    useState<ImageBundle[]>(startingBundles);
-  const [selectedImageToken, setSelectedImageToken] = useState<string | null>(
-    startingBundles.length > 0 && startingBundles[0].images.length > 0
-      ? startingBundles[0].images[0].mediaToken
-      : null,
-  );
   useImageEditCompleteEvent(async (event) => {
     const newBundle: ImageBundle = {
       images: event.edited_images.map(
@@ -49,12 +46,7 @@ export const HistoryStack = ({
       ),
     };
 
-    setImageBundles([...imageBundles, newBundle]);
-    const newlySelected = newBundle.images[0];
-    if (newlySelected) {
-      setSelectedImageToken(newlySelected.mediaToken);
-      onImageSelect(newlySelected);
-    }
+    onNewImageBundle(newBundle);
   });
 
   // This is used to force image reloads in different sessions
@@ -62,8 +54,6 @@ export const HistoryStack = ({
   const sessionRandBuster = useRef(Math.random());
 
   const handleClear = () => {
-    setImageBundles(startingBundles);
-    setSelectedImageToken(null);
     onClear();
   };
 
@@ -78,7 +68,7 @@ export const HistoryStack = ({
                 className={twMerge(
                   "relative aspect-square h-full w-full border-2 bg-transparent p-0 hover:bg-transparent hover:opacity-80",
                   selectedImageToken === image.mediaToken &&
-                    "border-primary hover:opacity-100",
+                  "border-primary hover:opacity-100",
                 )}
                 onClick={() => {
                   setSelectedImageToken(image.mediaToken);
diff --git a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx
index 4efcdcc732..83665ff922 100644
--- a/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx
+++ b/frontend/apps/editor3d/app/src/pages/PageEdit/PageEdit.tsx
@@ -27,7 +27,7 @@ const PAGE_ID: ModelPage = ModelPage.ImageEditor;
 const PageEdit = () => {
   //useStateSceneLoader();
 
-  const selectedImageModel : ImageModel | undefined = getSelectedImageModel(PAGE_ID);
+  const selectedImageModel: ImageModel | undefined = getSelectedImageModel(PAGE_ID);
 
   // State for canvas dimensions
   const canvasWidth = useRef<number>(1024);
@@ -46,6 +46,8 @@ const PageEdit = () => {
 
   // Use the Zustand store
   const store = useEditStore();
+  const addHistoryImageBundle = useEditStore((state) => state.addHistoryImageBundle);
+  const historyImageBundles = useEditStore((state) => state.historyImageBundles)
 
   // Pass store actions directly as callbacks
   useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
@@ -273,6 +275,7 @@ const PageEdit = () => {
       >
         <BaseImageSelector
           onImageSelect={(image: BaseSelectorImage) => {
+            addHistoryImageBundle({ images: [image] });
             store.setBaseImageInfo(image);
           }}
           showLoading={
@@ -305,11 +308,13 @@ const PageEdit = () => {
           onClear={() => {
             store.RESET();
           }}
-          startingBundles={[{ images: [store.baseImageInfo] } as ImageBundle]}
+          imageBundles={historyImageBundles}
           onImageSelect={(baseImage) => {
             store.clearLineNodes();
             store.setBaseImageInfo(baseImage);
           }}
+          onNewImageBundle={addHistoryImageBundle}
+          selectedImageToken={store.baseImageInfo?.mediaToken}
         />
       </div>
       <div
diff --git a/frontend/apps/editor3d/app/src/pages/PageEdit/stores/EditState.ts b/frontend/apps/editor3d/app/src/pages/PageEdit/stores/EditState.ts
index d6731c1b04..0a000df5d1 100644
--- a/frontend/apps/editor3d/app/src/pages/PageEdit/stores/EditState.ts
+++ b/frontend/apps/editor3d/app/src/pages/PageEdit/stores/EditState.ts
@@ -3,6 +3,7 @@ import { LineNode, generateId } from "~/pages/PageDraw/stores/SceneState";
 import { Node } from "~/pages/PageDraw/Node";
 import { BaseSelectorImage } from "../BaseImageSelector";
 import { FetchProxy } from "libs/tauri-utils/src/lib/FetchProxy";
+import { ImageBundle } from "../HistoryStack";
 
 export type ActiveEditTool = "select" | "edit" | "expand";
 export type EditOperation = "add" | "minus";
@@ -126,6 +127,11 @@ interface EditState {
   // pushHistory: () => void;
   // clearHistory: () => void;
   RESET: () => void;
+
+  // For the history stack
+  historyImageBundles: ImageBundle[];
+  clearHistoryImages: () => void;
+  addHistoryImageBundle: (bundle: ImageBundle) => void;
 }
 
 // Add a flag to prevent saving state during restoration
@@ -172,6 +178,9 @@ export const useEditStore = create<EditState>((set, get, store) => ({
   baseImageInfo: null,
   baseImageBitmap: null,
 
+  // History stack state
+  historyImageBundles: [],
+
   // Actions
   addNode: (node: Node) => {
     set((state) => {
@@ -1009,5 +1018,13 @@ export const useEditStore = create<EditState>((set, get, store) => ({
 
   RESET: () => {
     set(store.getInitialState());
-  }
+  },
+
+  clearHistoryImages: () => {
+    set({ historyImageBundles: [] });
+  },
+
+  addHistoryImageBundle(bundle) {
+    set((state) => ({ historyImageBundles: [...state.historyImageBundles, bundle] }));
+  },
 }));