From 1f39bffd4cc99161725e72946def5976608b759b Mon Sep 17 00:00:00 2001
From: Kerem Yilmaz <kerem@skyvern.com>
Date: Fri, 17 May 2024 11:03:00 -0700
Subject: [PATCH] Improve screenshot size & zoom (#340)

Co-authored-by: Muhammed Salih Altun <muhammedsalihaltun@gmail.com>
---
 .../src/components/ZoomableImage.tsx          | 64 +++++++++++++++++--
 .../src/routes/tasks/detail/StepArtifacts.tsx |  6 +-
 2 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/skyvern-frontend/src/components/ZoomableImage.tsx b/skyvern-frontend/src/components/ZoomableImage.tsx
index 3d0d3ebb..69fc679b 100644
--- a/skyvern-frontend/src/components/ZoomableImage.tsx
+++ b/skyvern-frontend/src/components/ZoomableImage.tsx
@@ -1,15 +1,65 @@
-import Zoom from "react-medium-image-zoom";
-import { AspectRatio } from "@/components/ui/aspect-ratio";
+import { useEffect, useState } from "react";
+import clsx from "clsx";
 
 type HTMLImageElementProps = React.ComponentProps<"img">;
 
 function ZoomableImage(props: HTMLImageElementProps) {
+  const [modalOpen, setModalOpen] = useState(false);
+  const [zoom, setZoom] = useState(false);
+
+  const openModal = () => {
+    setZoom(false);
+    setModalOpen(true);
+  };
+
+  const closeModal = () => {
+    setModalOpen(false);
+  };
+
+  useEffect(() => {
+    function handleEscKey(e: KeyboardEvent) {
+      if (modalOpen && e.key === "Escape") {
+        closeModal();
+      }
+    }
+    document.addEventListener("keydown", handleEscKey);
+    return () => {
+      document.removeEventListener("keydown", handleEscKey);
+    };
+  }, [modalOpen]);
+
   return (
-    <Zoom>
-      <AspectRatio ratio={16 / 9}>
-        <img {...props} />
-      </AspectRatio>
-    </Zoom>
+    <div>
+      <img {...props} onClick={openModal} />
+      {modalOpen && (
+        <div
+          className={clsx(
+            "fixed inset-0 z-50 flex justify-center bg-black bg-opacity-75 overflow-auto p-16",
+            {
+              "items-center": !zoom,
+              "items-baseline": zoom,
+            },
+          )}
+        >
+          <span
+            className="absolute top-4 right-4 text-white text-4xl cursor-pointer"
+            onClick={closeModal}
+          >
+            &times;
+          </span>
+          <img
+            {...props}
+            onClick={() => setZoom(!zoom)}
+            className={clsx({
+              "min-h-full object-contain h-full w-full m-0 cursor-zoom-in max-h-full max-w-full":
+                !zoom,
+              "min-h-full object-contain m-0 cursor-zoom-out max-w-none max-h-none mr-auto ml-auto":
+                zoom,
+            })}
+          />
+        </div>
+      )}
+    </div>
   );
 }
 
diff --git a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
index 8a62f6ab..d828b4e6 100644
--- a/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
+++ b/skyvern-frontend/src/routes/tasks/detail/StepArtifacts.tsx
@@ -111,7 +111,7 @@ function StepArtifacts({ id, stepProps }: Props) {
       </TabsContent>
       <TabsContent value="screenshot_llm">
         {llmScreenshots && llmScreenshots.length > 0 ? (
-          <div className="grid grid-cols-3 gap-4 p-4">
+          <div className="grid grid-cols-2 gap-4 p-4">
             {llmScreenshots.map((artifact, index) => (
               <ZoomableImage
                 key={index}
@@ -122,7 +122,7 @@ function StepArtifacts({ id, stepProps }: Props) {
             ))}
           </div>
         ) : isFetching ? (
-          <div className="grid grid-cols-3 gap-4 p-4">
+          <div className="grid grid-cols-2 gap-4 p-4">
             <Skeleton className="w-full h-full" />
             <Skeleton className="w-full h-full" />
             <Skeleton className="w-full h-full" />
@@ -133,7 +133,7 @@ function StepArtifacts({ id, stepProps }: Props) {
       </TabsContent>
       <TabsContent value="screenshot_action">
         {actionScreenshots && actionScreenshots.length > 0 ? (
-          <div className="grid grid-cols-3 gap-4 p-4">
+          <div className="grid grid-cols-2 gap-4 p-4">
             {actionScreenshots.map((artifact, index) => (
               <ZoomableImage
                 key={index}