Merge pull request #355 from abi/select-and-edit

Add "Select and Edit" functionality for easier edits

frontend/src/App.tsx

@@ -40,6 +40,8 @@ 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";
+import { useAppStore } from "./store/app-store";
 
 const IS_OPENAI_DOWN = false;
 
@@ -54,6 +56,8 @@ function App() {
   const [updateInstruction, setUpdateInstruction] = useState("");
   const [isImportedFromCode, setIsImportedFromCode] = useState<boolean>(false);
 
+  const { disableInSelectAndEditMode } = useAppStore();
+
   // Settings
   const [settings, setSettings] = usePersistedState<Settings>(
     {
@@ -94,6 +98,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);
 
@@ -149,6 +157,7 @@ function App() {
     setAppHistory([]);
     setCurrentVersion(null);
     setShouldIncludeResultImage(false);
+    disableInSelectAndEditMode();
   };
 
   const regenerate = () => {
@@ -237,7 +246,9 @@ function App() {
                 parentIndex: parentVersion,
                 code,
                 inputs: {
-                  prompt: updateInstruction,
+                  prompt: params.history
+                    ? params.history[params.history.length - 1]
+                    : updateInstruction,
                 },
               },
             ];
@@ -279,7 +290,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."
@@ -297,7 +311,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();
@@ -379,10 +403,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} />
@@ -485,7 +506,7 @@ function App() {
                       />
                     </div>
                     <Button
-                      onClick={doUpdate}
+                      onClick={() => doUpdate(updateInstruction)}
                       className="dark:text-white dark:bg-gray-700 update-btn"
                     >
                       Update
@@ -498,6 +519,9 @@ function App() {
                     >
                       🔄 Regenerate
                     </Button>
+                    {showSelectAndEditFeature && (
+                      <SelectAndEditModeToggleButton />
+                    )}
                   </div>
                   <div className="flex justify-end items-center mt-2">
                     <TipLink />
@@ -626,10 +650,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

frontend/src/components/Preview.tsx

@@ -1,21 +1,35 @@
-import { useEffect, useRef } from "react";
+import { useEffect, useRef, useState } from "react";
 import classNames from "classnames";
 import useThrottle from "../hooks/useThrottle";
+import EditPopup from "./select-and-edit/EditPopup";
 
 interface Props {
   code: string;
   device: "mobile" | "desktop";
+  doUpdate: (updateInstruction: string, selectedElement?: HTMLElement) => void;
 }
 
-function Preview({ code, device }: Props) {
+function Preview({ code, device, doUpdate }: Props) {
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
 
   // Don't update code more often than every 200ms.
   const throttledCode = useThrottle(code, 200);
 
+  // Select and edit functionality
+  const [clickEvent, setClickEvent] = useState<MouseEvent | null>(null);
+
   useEffect(() => {
-    if (iframeRef.current) {
-      iframeRef.current.srcdoc = throttledCode;
+    const iframe = iframeRef.current;
+    if (iframe) {
+      iframe.srcdoc = throttledCode;
+
+      // Set up click handler for select and edit funtionality
+      iframe.addEventListener("load", function () {
+        iframe.contentWindow?.document.body.addEventListener(
+          "click",
+          setClickEvent
+        );
+      });
     }
   }, [throttledCode]);
 
@@ -34,6 +48,7 @@ function Preview({ code, device }: Props) {
           }
         )}
       ></iframe>
+      <EditPopup event={clickEvent} iframeRef={iframeRef} doUpdate={doUpdate} />
     </div>
   );
 }

frontend/src/components/select-and-edit/EditPopup.tsx

@@ -0,0 +1,143 @@
+import React, { useEffect, useRef, useState } from "react";
+import { Textarea } from "../ui/textarea";
+import { Button } from "../ui/button";
+import { addHighlight, getAdjustedCoordinates, removeHighlight } from "./utils";
+import { useAppStore } from "../../store/app-store";
+
+interface EditPopupProps {
+  event: MouseEvent | null;
+  iframeRef: React.RefObject<HTMLIFrameElement>;
+  doUpdate: (updateInstruction: string, selectedElement?: HTMLElement) => void;
+}
+
+const EditPopup: React.FC<EditPopupProps> = ({
+  event,
+  iframeRef,
+  doUpdate,
+}) => {
+  // App state
+  const { inSelectAndEditMode } = useAppStore();
+
+  // Create a wrapper ref to store inSelectAndEditMode so the value is not stale
+  // in a event listener
+  const inSelectAndEditModeRef = useRef(inSelectAndEditMode);
+
+  // Update the ref whenever the state changes
+  useEffect(() => {
+    inSelectAndEditModeRef.current = inSelectAndEditMode;
+  }, [inSelectAndEditMode]);
+
+  // Popup state
+  const [popupVisible, setPopupVisible] = useState(false);
+  const [popupPosition, setPopupPosition] = useState({ x: 0, y: 0 });
+
+  // Edit state
+  const [selectedElement, setSelectedElement] = useState<
+    HTMLElement | undefined
+  >(undefined);
+  const [updateText, setUpdateText] = useState("");
+
+  // Textarea ref for focusing
+  const textareaRef = useRef<HTMLTextAreaElement | null>(null);
+
+  function onUpdate(updateText: string) {
+    // Perform the update
+    doUpdate(
+      updateText,
+      selectedElement ? removeHighlight(selectedElement) : selectedElement
+    );
+
+    // Unselect the element
+    setSelectedElement(undefined);
+
+    // Hide the popup
+    setPopupVisible(false);
+  }
+
+  // Remove highlight and reset state when not in select and edit mode
+  useEffect(() => {
+    if (!inSelectAndEditMode) {
+      if (selectedElement) removeHighlight(selectedElement);
+      setSelectedElement(undefined);
+      setPopupVisible(false);
+    }
+  }, [inSelectAndEditMode, selectedElement]);
+
+  // Handle the click event
+  useEffect(() => {
+    // Return if not in select and edit mode
+    if (!inSelectAndEditModeRef.current || !event) {
+      return;
+    }
+
+    // Prevent default to avoid issues like label clicks triggering textareas, etc.
+    event.preventDefault();
+
+    const targetElement = event.target as HTMLElement;
+
+    // Return if no target element
+    if (!targetElement) return;
+
+    // Highlight and set the selected element
+    setSelectedElement((prev) => {
+      // Remove style from previous element
+      if (prev) {
+        removeHighlight(prev);
+      }
+      return addHighlight(targetElement);
+    });
+
+    // Calculate adjusted coordinates
+    const adjustedCoordinates = getAdjustedCoordinates(
+      event.clientX,
+      event.clientY,
+      iframeRef.current?.getBoundingClientRect()
+    );
+
+    // Show the popup at the click position
+    setPopupVisible(true);
+    setPopupPosition({ x: adjustedCoordinates.x, y: adjustedCoordinates.y });
+
+    // Reset the update text
+    setUpdateText("");
+
+    // Focus the textarea
+    textareaRef.current?.focus();
+  }, [event, iframeRef]);
+
+  // Focus the textarea when the popup is visible (we can't do this only when handling the click event
+  // because the textarea is not rendered yet)
+  // We need to also do it in the click event because popupVisible doesn't change values in that event
+  useEffect(() => {
+    if (popupVisible) {
+      textareaRef.current?.focus();
+    }
+  }, [popupVisible]);
+
+  if (!popupVisible) return;
+
+  return (
+    <div
+      className="absolute bg-white p-4 border border-gray-300 rounded shadow-lg w-60"
+      style={{ top: popupPosition.y, left: popupPosition.x }}
+    >
+      <Textarea
+        ref={textareaRef}
+        value={updateText}
+        onChange={(e) => setUpdateText(e.target.value)}
+        placeholder="Tell the AI what to change about this element..."
+        onKeyDown={(e) => {
+          if (e.key === "Enter") {
+            e.preventDefault();
+            onUpdate(updateText);
+          }
+        }}
+      />
+      <div className="flex justify-end mt-2">
+        <Button onClick={() => onUpdate(updateText)}>Update</Button>
+      </div>
+    </div>
+  );
+};
+
+export default EditPopup;

frontend/src/components/select-and-edit/SelectAndEditModeToggleButton.tsx

@@ -0,0 +1,22 @@
+import { GiClick } from "react-icons/gi";
+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"}
+    >
+      <GiClick className="text-lg" />
+      <span>
+        {inSelectAndEditMode ? "Exit selection mode" : "Select and update"}
+      </span>
+    </Button>
+  );
+}
+
+export default SelectAndEditModeToggleButton;

frontend/src/components/select-and-edit/utils.ts

@@ -0,0 +1,22 @@
+export function removeHighlight(element: HTMLElement) {
+  element.style.outline = "";
+  element.style.backgroundColor = "";
+  return element;
+}
+
+export function addHighlight(element: HTMLElement) {
+  element.style.outline = "2px dashed #1846db";
+  element.style.backgroundColor = "#bfcbf5";
+  return element;
+}
+
+export function getAdjustedCoordinates(
+  x: number,
+  y: number,
+  rect: DOMRect | undefined
+) {
+  const offsetX = rect ? rect.left : 0;
+  const offsetY = rect ? rect.top : 0;
+
+  return { x: x + offsetX, y: y + offsetY };
+}

frontend/src/store/app-store.ts

@@ -2,9 +2,14 @@ import { create } from "zustand";
 
 // Store for app-wide state
 interface AppStore {
-  inputMode: "image" | "video";
+  inSelectAndEditMode: boolean;
+  toggleInSelectAndEditMode: () => void;
+  disableInSelectAndEditMode: () => void;
 }
 
-export const useStore = create<AppStore>(() => ({
-  inputMode: "image",
+export const useAppStore = create<AppStore>((set) => ({
+  inSelectAndEditMode: false,
+  toggleInSelectAndEditMode: () =>
+    set((state) => ({ inSelectAndEditMode: !state.inSelectAndEditMode })),
+  disableInSelectAndEditMode: () => set({ inSelectAndEditMode: false }),
 }));