clean up more of the code

frontend/src/components/Preview.tsx

@@ -6,29 +6,38 @@ import EditPopup from "./select-and-edit/EditPopup";
 interface Props {
   code: string;
   device: "mobile" | "desktop";
+  doUpdate: (
+    selectedUpdateInstruction?: string,
+    selectedElement?: HTMLElement
+  ) => void;
+  inSelectAndEditMode: boolean;
 }
 
-function Preview({ code, device }: Props) {
+function Preview({ code, device, doUpdate, inSelectAndEditMode }: Props) {
   const iframeRef = useRef<HTMLIFrameElement | null>(null);
 
   // Don't update code more often than every 200ms.
   const throttledCode = useThrottle(code, 200);
 
-  const [selectedElement, setSelectedElement] = useState<HTMLElement | null>(
+  const inSelectAndEditModeRef = useRef(inSelectAndEditMode); // Create a ref for the state
-    null
+
-  );
+  const [selectedElement, setSelectedElement] = useState<
+    HTMLElement | undefined
+  >(undefined);
   const [popupVisible, setPopupVisible] = useState(false);
   const [popupPosition, setPopupPosition] = useState({ x: 0, y: 0 });
-  const [editText, setEditText] = useState("");
 
-  console.log(selectedElement);
+  function removeHighlight(element: HTMLElement) {
+    element.style.outline = "";
+    element.style.backgroundColor = "";
+    return element;
+  }
 
   function updateHighlight(targetElement: HTMLElement) {
     setSelectedElement((prev) => {
       // Remove style from previous element
       if (prev) {
-        prev.style.outline = "";
+        removeHighlight(prev);
-        prev.style.backgroundColor = "";
       }
       // Add style to new element
       targetElement.style.outline = "2px dashed #1846db";
@@ -38,6 +47,11 @@ function Preview({ code, device }: Props) {
   }
 
   function handleClick(event: MouseEvent) {
+    // Return if not in select and edit mode
+    if (!inSelectAndEditModeRef.current) {
+      return;
+    }
+
     const { clientX, clientY } = event;
 
     // Prevent default to avoid issues like label clicks triggering textareas, etc.
@@ -70,13 +84,28 @@ function Preview({ code, device }: Props) {
     setPopupVisible(true);
   }
 
-  function handleEditSubmit() {
+  function handleEditSubmit(editText: string) {
-    if (selectedElement) {
+    doUpdate(
-      selectedElement.innerText = editText;
+      editText,
-    }
+      selectedElement ? removeHighlight(selectedElement) : selectedElement
+    );
+    setSelectedElement(undefined);
     setPopupVisible(false);
   }
 
+  useEffect(() => {
+    if (!inSelectAndEditMode) {
+      if (selectedElement) removeHighlight(selectedElement);
+      setSelectedElement(undefined);
+      setPopupVisible(false);
+    }
+  }, [inSelectAndEditMode, selectedElement]);
+
+  // Update the ref whenever the state changes
+  useEffect(() => {
+    inSelectAndEditModeRef.current = inSelectAndEditMode;
+  }, [inSelectAndEditMode]);
+
   useEffect(() => {
     const iframe = iframeRef.current;
     if (iframe) {
@@ -111,8 +140,6 @@ function Preview({ code, device }: Props) {
         {...{
           popupVisible,
           popupPosition,
-          editText,
-          setEditText,
           handleEditSubmit,
         }}
       />

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

@@ -1,22 +1,19 @@
-import React, { useEffect, useRef } from "react";
+import React, { useEffect, useRef, useState } from "react";
 import { Textarea } from "../ui/textarea";
 import { Button } from "../ui/button";
 
 interface EditPopupProps {
   popupVisible: boolean;
   popupPosition: { x: number; y: number };
-  editText: string;
+  handleEditSubmit: (editText: string) => void;
-  setEditText: (text: string) => void;
-  handleEditSubmit: () => void;
 }
 
 const EditPopup: React.FC<EditPopupProps> = ({
   popupVisible,
   popupPosition,
-  editText,
-  setEditText,
   handleEditSubmit,
 }) => {
+  const [editText, setEditText] = useState("");
   const textareaRef = useRef<HTMLTextAreaElement | null>(null);
 
   useEffect(() => {
@@ -45,7 +42,7 @@ const EditPopup: React.FC<EditPopupProps> = ({
         placeholder="Tell the AI what to change about this element..."
       />
       <div className="flex justify-end mt-2">
-        <Button onClick={handleEditSubmit}>Update</Button>
+        <Button onClick={() => handleEditSubmit(editText)}>Update</Button>
       </div>
     </div>
   );