From 3f4cf2895eeb8d3c29858ea3c6dfd7e91770a90a Mon Sep 17 00:00:00 2001
From: Abi Raja <abimanyuraja@gmail.com>
Date: Tue, 14 Nov 2023 15:31:24 -0500
Subject: [PATCH] better UX for different app states

---
 backend/main.py              |  2 ++
 frontend/src/App.tsx         | 31 ++++++++++++++++++++++---------
 frontend/src/generateCode.ts |  5 ++++-
 3 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/backend/main.py b/backend/main.py
index ad8d02c..e03cbc6 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -46,3 +46,5 @@ async def stream_code_test(websocket: WebSocket):
 
     # Write the messages dict into a log so that we can debug later
     write_logs(prompt_messages, completion)
+
+    await websocket.close()
diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index 6932925..fcefcb8 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -5,31 +5,44 @@ import Preview from "./components/Preview";
 import { generateCode } from "./generateCode";
 
 function App() {
+  const [appState, setAppState] = useState<"INITIAL" | "CODING" | "CODE_READY">(
+    "INITIAL"
+  );
   const [generatedCode, setGeneratedCode] = useState<string>("");
   const [referenceImages, setReferenceImages] = useState<string[]>([]);
 
   function startCodeGeneration(referenceImages: string[]) {
+    setAppState("CODING");
     setReferenceImages(referenceImages);
-    generateCode(referenceImages[0], function (token) {
-      setGeneratedCode((prev) => prev + token);
-    });
+    generateCode(
+      referenceImages[0],
+      function (token) {
+        setGeneratedCode((prev) => prev + token);
+      },
+      function () {
+        setAppState("CODE_READY");
+      }
+    );
   }
 
   return (
     <div className="mx-auto mt-10 max-w-[1000px]">
       <h1 className="text-2xl mb-4">Drag & Drop a Screenshot</h1>
-      {referenceImages.length > 0 && (
-        <img className="w-[300px]" src={referenceImages[0]} alt="Reference" />
-      )}
 
-      {referenceImages.length === 0 && (
+      {appState === "INITIAL" && (
         <>
           <ImageUpload setReferenceImages={startCodeGeneration} />
         </>
       )}
 
-      <CodePreview code={generatedCode} />
-      <Preview code={generatedCode} />
+      {(appState === "CODING" || appState === "CODE_READY") && (
+        <>
+          <img className="w-[300px]" src={referenceImages[0]} alt="Reference" />
+          {/* Show code preview only when coding */}
+          {appState === "CODING" && <CodePreview code={generatedCode} />}
+          <Preview code={generatedCode} />
+        </>
+      )}
     </div>
   );
 }
diff --git a/frontend/src/generateCode.ts b/frontend/src/generateCode.ts
index 6420b3b..eb5bbf4 100644
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@@ -4,7 +4,8 @@ const WS_BACKEND_URL = import.meta.env.VITE_WS_BACKEND_URL;
 
 export function generateCode(
   imageUrl: string,
-  onChange: (chunk: string) => void
+  onChange: (chunk: string) => void,
+  onComplete: () => void
 ) {
   const wsUrl = `${WS_BACKEND_URL}/generate-code`;
   console.log("Connecting to backend @ ", wsUrl);
@@ -33,6 +34,8 @@ export function generateCode(
       toast.error(
         "We ran into an error. Try again or contact support at support@picoapps.xyz"
       );
+    } else {
+      onComplete();
     }
   });