From 53ae82f0b9ba67cd4b075035783ae50db92a8bd8 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Fri, 23 Feb 2024 12:54:15 -0300
Subject: [PATCH] Add VertexBuildTypeAPI import and refactor handleBuildUpdate
 function

---
 src/frontend/src/stores/flowStore.ts | 37 +++++++++++++++++-----------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/frontend/src/stores/flowStore.ts b/src/frontend/src/stores/flowStore.ts
index 588a8ae21..549b0a90c 100644
--- a/src/frontend/src/stores/flowStore.ts
+++ b/src/frontend/src/stores/flowStore.ts
@@ -11,6 +11,7 @@ import {
 import { create } from "zustand";
 import { BuildStatus } from "../constants/enums";
 import { getFlowPool, updateFlowInDatabase } from "../controllers/API";
+import { VertexBuildTypeAPI } from "../types/api";
 import {
   NodeDataType,
   NodeType,
@@ -151,7 +152,6 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         ? change(get().nodes.find((node) => node.id === id)!)
         : change;
 
-    console.log(newChange);
     get().setNodes((oldNodes) =>
       oldNodes.map((node) => {
         if (node.id === id) {
@@ -375,17 +375,15 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     const setSuccessData = useAlertStore.getState().setSuccessData;
     const setErrorData = useAlertStore.getState().setErrorData;
     const setNoticeData = useAlertStore.getState().setNoticeData;
-    function handleBuildUpdate(data: any) {
-      const responseData = data.data[data.id];
-      // responseData MAY contain inactive_vertices or it
-      // may be null, so we need to check for it
-      if (responseData && responseData.inactive_vertices) {
-        useFlowStore
-          .getState()
-          .setInactiveNodes(responseData.inactive_vertices);
+    function handleBuildUpdate(
+      vertexBuildData: VertexBuildTypeAPI,
+      status: BuildStatus
+    ) {
+      if (vertexBuildData && vertexBuildData.inactive_vertices) {
+        get().removeFromVerticesBuild(vertexBuildData.inactive_vertices);
       }
-      get().addDataToFlowPool(data.data[data.id], data.id);
-      useFlowStore.getState().updateBuildStatus([data.id], BuildStatus.BUILT);
+      get().addDataToFlowPool(vertexBuildData, vertexBuildData.id);
+      useFlowStore.getState().updateBuildStatus([vertexBuildData.id], status);
     }
     await updateFlowInDatabase({
       data: {
@@ -418,6 +416,7 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
         useFlowStore.getState().updateBuildStatus(idList, BuildStatus.BUILDING);
       },
     });
+    get().revertAllVerticesToBuild();
   },
   getFlow: () => {
     return {
@@ -439,10 +438,20 @@ const useFlowStore = create<FlowStoreType>((set, get) => ({
     set({ verticesBuild: vertices });
   },
   verticesBuild: [],
-  setInactiveNodes(newState) {
-    set({ inactiveNodes: newState });
+  revertAllVerticesToBuild: () => {
+    // set all vertices to TO_BUILD
+    const verticesIds = get()
+      .nodes.filter((node) => node.data.build_status === BuildStatus.BUILDING)
+      .map((node) => node.id);
+    get().updateBuildStatus(verticesIds, BuildStatus.TO_BUILD);
+  },
+  removeFromVerticesBuild: (vertices: string[]) => {
+    set({
+      verticesBuild: get().verticesBuild.filter(
+        (vertex) => !vertices.includes(vertex)
+      ),
+    });
   },
-  inactiveNodes: [],
 }));
 
 export default useFlowStore;