diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/Preview.tsx
index c643e69..2a25ac1 100644
--- a/frontend/src/components/Preview.tsx
+++ b/frontend/src/components/Preview.tsx
@@ -19,103 +19,19 @@ function Preview({ code, device, doUpdate, inSelectAndEditMode }: Props) {
   // Don't update code more often than every 200ms.
   const throttledCode = useThrottle(code, 200);
 
-  const inSelectAndEditModeRef = useRef(inSelectAndEditMode); // Create a ref for the state
-
-  const [selectedElement, setSelectedElement] = useState<
-    HTMLElement | undefined
-  >(undefined);
-  const [popupVisible, setPopupVisible] = useState(false);
-  const [popupPosition, setPopupPosition] = useState({ x: 0, y: 0 });
-
-  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) {
-        removeHighlight(prev);
-      }
-      // Add style to new element
-      targetElement.style.outline = "2px dashed #1846db";
-      targetElement.style.backgroundColor = "#bfcbf5";
-      return targetElement;
-    });
-  }
-
-  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.
-    event.preventDefault();
-
-    const targetElement = event.target as HTMLElement;
-
-    // Return if no target element
-    if (!targetElement) {
-      return;
-    }
-
-    // Highlight the selected element
-    updateHighlight(targetElement);
-
-    // Show popup at click position, slightly offset to be right under the cursor
-    setPopupVisible(false);
-
-    // Calculate offsets
-    const rect = iframeRef.current?.getBoundingClientRect();
-    const offsetX = rect ? rect.left : 0;
-    const offsetY = rect ? rect.top : 0;
-
-    // Adjust for scale
-    const scale = 1; // the scale factor applied to the iframe
-    const scaledX = clientX / scale + offsetX;
-    const scaledY = clientY / scale + offsetY;
-
-    setPopupPosition({ x: scaledX, y: scaledY });
-    setPopupVisible(true);
-  }
-
-  function handleEditSubmit(editText: string) {
-    doUpdate(
-      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]);
+  // Select and edit functionality
+  const [clickEvent, setClickEvent] = useState<MouseEvent | null>(null);
 
   useEffect(() => {
     const iframe = iframeRef.current;
     if (iframe) {
       iframe.srcdoc = throttledCode;
 
-      // Related to
+      // Set up click handler for select and edit funtionality
       iframe.addEventListener("load", function () {
         iframe.contentWindow?.document.body.addEventListener(
           "click",
-          handleClick
+          setClickEvent
         );
       });
     }
@@ -137,11 +53,10 @@ function Preview({ code, device, doUpdate, inSelectAndEditMode }: Props) {
         )}
       ></iframe>
       <EditPopup
-        {...{
-          popupVisible,
-          popupPosition,
-          handleEditSubmit,
-        }}
+        event={clickEvent}
+        iframeRef={iframeRef}
+        inSelectAndEditMode={inSelectAndEditMode}
+        doUpdate={doUpdate}
       />
     </div>
   );
diff --git a/frontend/src/components/select-and-edit/EditPopup.tsx b/frontend/src/components/select-and-edit/EditPopup.tsx
index a833f00..de05ae7 100644
--- a/frontend/src/components/select-and-edit/EditPopup.tsx
+++ b/frontend/src/components/select-and-edit/EditPopup.tsx
@@ -2,20 +2,60 @@ import React, { useEffect, useRef, useState } from "react";
 import { Textarea } from "../ui/textarea";
 import { Button } from "../ui/button";
 
+function removeHighlight(element: HTMLElement) {
+  element.style.outline = "";
+  element.style.backgroundColor = "";
+  return element;
+}
+
 interface EditPopupProps {
-  popupVisible: boolean;
-  popupPosition: { x: number; y: number };
-  handleEditSubmit: (editText: string) => void;
+  event: MouseEvent | null;
+  inSelectAndEditMode: boolean;
+  doUpdate: (
+    selectedUpdateInstruction?: string,
+    selectedElement?: HTMLElement
+  ) => void;
+  iframeRef: React.RefObject<HTMLIFrameElement>;
 }
 
 const EditPopup: React.FC<EditPopupProps> = ({
-  popupVisible,
-  popupPosition,
-  handleEditSubmit,
+  event,
+  inSelectAndEditMode,
+  doUpdate,
+  iframeRef,
 }) => {
   const [editText, setEditText] = useState("");
+  const [selectedElement, setSelectedElement] = useState<
+    HTMLElement | undefined
+  >(undefined);
+  const [popupVisible, setPopupVisible] = useState(false);
+  const [popupPosition, setPopupPosition] = useState({ x: 0, y: 0 });
+
+  const inSelectAndEditModeRef = useRef(inSelectAndEditMode); // Create a ref for the state
   const textareaRef = useRef<HTMLTextAreaElement | null>(null);
 
+  function updateHighlight(targetElement: HTMLElement) {
+    setSelectedElement((prev) => {
+      // Remove style from previous element
+      if (prev) {
+        removeHighlight(prev);
+      }
+      // Add style to new element
+      targetElement.style.outline = "2px dashed #1846db";
+      targetElement.style.backgroundColor = "#bfcbf5";
+      return targetElement;
+    });
+  }
+
+  function handleEditSubmit(editText: string) {
+    doUpdate(
+      editText,
+      selectedElement ? removeHighlight(selectedElement) : selectedElement
+    );
+    setSelectedElement(undefined);
+    setPopupVisible(false);
+  }
+
   useEffect(() => {
     if (popupVisible && textareaRef.current) {
       textareaRef.current.focus();
@@ -28,6 +68,57 @@ const EditPopup: React.FC<EditPopupProps> = ({
     }
   }, [popupVisible, setEditText]);
 
+  useEffect(() => {
+    if (!inSelectAndEditMode) {
+      if (selectedElement) removeHighlight(selectedElement);
+      setSelectedElement(undefined);
+      setPopupVisible(false);
+    }
+  }, [inSelectAndEditMode, selectedElement]);
+
+  useEffect(() => {
+    // Return if not in select and edit mode
+    if (!inSelectAndEditModeRef.current || !event) {
+      return;
+    }
+
+    const { clientX, clientY } = event;
+
+    // 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 the selected element
+    updateHighlight(targetElement);
+
+    // Show popup at click position, slightly offset to be right under the cursor
+    setPopupVisible(false);
+
+    // Calculate offsets
+    const rect = iframeRef.current?.getBoundingClientRect();
+    const offsetX = rect ? rect.left : 0;
+    const offsetY = rect ? rect.top : 0;
+
+    // Adjust for scale
+    const scale = 1; // the scale factor applied to the iframe
+    const scaledX = clientX / scale + offsetX;
+    const scaledY = clientY / scale + offsetY;
+
+    setPopupPosition({ x: scaledX, y: scaledY });
+    setPopupVisible(true);
+  }, [event]);
+
+  // Update the ref whenever the state changes
+  useEffect(() => {
+    inSelectAndEditModeRef.current = inSelectAndEditMode;
+  }, [inSelectAndEditMode]);
+
   if (!popupVisible) return;
 
   return (