From 9cfcf928d0a6c858f06ba845ed4639896f473b34 Mon Sep 17 00:00:00 2001
From: Abi Raja <abimanyuraja@gmail.com>
Date: Thu, 30 May 2024 09:51:44 -0400
Subject: [PATCH] isolate functionality

use zustand to share state rather than prop drilling

isolate more functionality away into a separate component
---
 frontend/src/App.tsx                          | 48 +++++++++++++++----
 frontend/src/components/Preview.tsx           | 11 ++---
 .../components/select-and-edit/EditPopup.tsx  |  5 +-
 .../SelectAndEditModeToggleButton.tsx         | 18 +++++++
 frontend/src/store/app-store.ts               |  7 ++-
 5 files changed, 68 insertions(+), 21 deletions(-)
 create mode 100644 frontend/src/components/select-and-edit/SelectAndEditModeToggleButton.tsx

diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 1cb02e1..77785db 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -40,6 +40,7 @@ import ModelSettingsSection from "./components/ModelSettingsSection";
 import { extractHtml } from "./components/preview/extractHtml";
 import useBrowserTabIndicator from "./hooks/useBrowserTabIndicator";
 import TipLink from "./components/core/TipLink";
+import SelectAndEditModeToggleButton from "./components/select-and-edit/SelectAndEditModeToggleButton";
 
 const IS_OPENAI_DOWN = false;
 
@@ -93,6 +94,10 @@ function App() {
     selectedCodeGenerationModel !== CodeGenerationModel.GPT_4O_2024_05_13 &&
     appState === AppState.INITIAL;
 
+  const showSelectAndEditFeature =
+    selectedCodeGenerationModel === CodeGenerationModel.GPT_4O_2024_05_13 &&
+    settings.generatedCodeConfig === Stack.HTML_TAILWIND;
+
   // Indicate coding state using the browser tab's favicon and title
   useBrowserTabIndicator(appState === AppState.CODING);
 
@@ -236,7 +241,9 @@ function App() {
                 parentIndex: parentVersion,
                 code,
                 inputs: {
-                  prompt: updateInstruction,
+                  prompt: params.history
+                    ? params.history[params.history.length - 1]
+                    : updateInstruction,
                 },
               },
             ];
@@ -278,7 +285,10 @@ function App() {
   }
 
   // Subsequent updates
-  async function doUpdate() {
+  async function doUpdate(
+    updateInstruction: string,
+    selectedElement?: HTMLElement
+  ) {
     if (currentVersion === null) {
       toast.error(
         "No current version set. Contact support or open a Github issue."
@@ -296,7 +306,17 @@ function App() {
       return;
     }
 
-    const updatedHistory = [...historyTree, updateInstruction];
+    let modifiedUpdateInstruction = updateInstruction;
+
+    // Send in a reference to the selected element if it exists
+    if (selectedElement) {
+      modifiedUpdateInstruction =
+        updateInstruction +
+        " referring to this element specifically: " +
+        selectedElement.outerHTML;
+    }
+
+    const updatedHistory = [...historyTree, modifiedUpdateInstruction];
 
     if (shouldIncludeResultImage) {
       const resultImage = await takeScreenshot();
@@ -378,10 +398,7 @@ function App() {
         />
       )}
       <div className="lg:fixed lg:inset-y-0 lg:z-40 lg:flex lg:w-96 lg:flex-col">
-        <div
-          className="flex grow flex-col gap-y-2 overflow-y-auto border-r 
-          border-gray-200 bg-white px-6 dark:bg-zinc-950 dark:text-white"
-        >
+        <div className="flex grow flex-col gap-y-2 overflow-y-auto border-r border-gray-200 bg-white px-6 dark:bg-zinc-950 dark:text-white">
           <div className="flex items-center justify-between mt-10 mb-2">
             <h1 className="text-2xl ">Screenshot to Code</h1>
             <SettingsDialog settings={settings} setSettings={setSettings} />
@@ -484,7 +501,7 @@ function App() {
                       />
                     </div>
                     <Button
-                      onClick={doUpdate}
+                      onClick={() => doUpdate(updateInstruction)}
                       className="dark:text-white dark:bg-gray-700 update-btn"
                     >
                       Update
@@ -497,6 +514,9 @@ function App() {
                     >
                       🔄 Regenerate
                     </Button>
+                    {showSelectAndEditFeature && (
+                      <SelectAndEditModeToggleButton />
+                    )}
                   </div>
                   <div className="flex justify-end items-center mt-2">
                     <TipLink />
@@ -625,10 +645,18 @@ function App() {
                 </div>
               </div>
               <TabsContent value="desktop">
-                <Preview code={previewCode} device="desktop" />
+                <Preview
+                  code={previewCode}
+                  device="desktop"
+                  doUpdate={doUpdate}
+                />
               </TabsContent>
               <TabsContent value="mobile">
-                <Preview code={previewCode} device="mobile" />
+                <Preview
+                  code={previewCode}
+                  device="mobile"
+                  doUpdate={doUpdate}
+                />
               </TabsContent>
               <TabsContent value="code">
                 <CodeTab
diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/Preview.tsx
index 2a25ac1..8dc2f22 100644
--- a/frontend/src/components/Preview.tsx
+++ b/frontend/src/components/Preview.tsx
@@ -2,18 +2,17 @@ import { useEffect, useRef, useState } from "react";
 import classNames from "classnames";
 import useThrottle from "../hooks/useThrottle";
 import EditPopup from "./select-and-edit/EditPopup";
+import { useAppStore } from "../store/app-store";
 
 interface Props {
   code: string;
   device: "mobile" | "desktop";
-  doUpdate: (
-    selectedUpdateInstruction?: string,
-    selectedElement?: HTMLElement
-  ) => void;
-  inSelectAndEditMode: boolean;
+  doUpdate: (updateInstruction: string, selectedElement?: HTMLElement) => void;
 }
 
-function Preview({ code, device, doUpdate, inSelectAndEditMode }: Props) {
+function Preview({ code, device, doUpdate }: Props) {
+  const { inSelectAndEditMode } = useAppStore();
+
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
 
   // Don't update code more often than every 200ms.
diff --git a/frontend/src/components/select-and-edit/EditPopup.tsx b/frontend/src/components/select-and-edit/EditPopup.tsx
index b9738ec..5bc9d42 100644
--- a/frontend/src/components/select-and-edit/EditPopup.tsx
+++ b/frontend/src/components/select-and-edit/EditPopup.tsx
@@ -6,10 +6,7 @@ import { addHighlight, getAdjustedCoordinates, removeHighlight } from "./utils";
 interface EditPopupProps {
   event: MouseEvent | null;
   inSelectAndEditMode: boolean;
-  doUpdate: (
-    selectedUpdateInstruction?: string,
-    selectedElement?: HTMLElement
-  ) => void;
+  doUpdate: (updateInstruction: string, selectedElement?: HTMLElement) => void;
   iframeRef: React.RefObject<HTMLIFrameElement>;
 }
 
diff --git a/frontend/src/components/select-and-edit/SelectAndEditModeToggleButton.tsx b/frontend/src/components/select-and-edit/SelectAndEditModeToggleButton.tsx
new file mode 100644
index 0000000..f076c9b
--- /dev/null
+++ b/frontend/src/components/select-and-edit/SelectAndEditModeToggleButton.tsx
@@ -0,0 +1,18 @@
+import { useAppStore } from "../../store/app-store";
+import { Button } from "../ui/button";
+
+function SelectAndEditModeToggleButton() {
+  const { inSelectAndEditMode, toggleInSelectAndEditMode } = useAppStore();
+
+  return (
+    <Button
+      onClick={toggleInSelectAndEditMode}
+      className="flex items-center gap-x-2 dark:text-white dark:bg-gray-700 regenerate-btn"
+      variant={inSelectAndEditMode ? "destructive" : "default"}
+    >
+      {inSelectAndEditMode ? "Exit selection mode" : "🖱️ Select and Edit"}
+    </Button>
+  );
+}
+
+export default SelectAndEditModeToggleButton;
diff --git a/frontend/src/store/app-store.ts b/frontend/src/store/app-store.ts
index 7298c72..26b7994 100644
--- a/frontend/src/store/app-store.ts
+++ b/frontend/src/store/app-store.ts
@@ -2,9 +2,14 @@ import { create } from "zustand";
 
 // Store for app-wide state
 interface AppStore {
+  inSelectAndEditMode: boolean;
   inputMode: "image" | "video";
+  toggleInSelectAndEditMode: () => void;
 }
 
-export const useStore = create<AppStore>(() => ({
+export const useAppStore = create<AppStore>((set) => ({
   inputMode: "image",
+  inSelectAndEditMode: false,
+  toggleInSelectAndEditMode: () =>
+    set((state) => ({ inSelectAndEditMode: !state.inSelectAndEditMode })),
 }));