From f81c1c415211d01c24920d4f847e82d6a2fb2397 Mon Sep 17 00:00:00 2001
From: Hassan El Mghari <hassan4709@gmail.com>
Date: Sat, 4 Mar 2023 01:22:00 -0500
Subject: [PATCH] tracking room IDs in localstorage for future features

---
 .example.env          |  3 +++
 pages/api/generate.ts | 29 +++++++++++++++++++++--------
 pages/dream.tsx       | 12 +++++++++---
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/.example.env b/.example.env
index 3b812e93..b4017c06 100644
--- a/.example.env
+++ b/.example.env
@@ -3,3 +3,6 @@ REPLICATE_API_KEY=
 # Optional, if you're doing rate limiting
 UPSTASH_REDIS_REST_URL=
 UPSTASH_REDIS_REST_TOKEN=
+
+# Optional, if you're planning to use Upload.io (you need to if you have 100+ uploads)
+NEXT_PUBLIC_UPLOAD_API_KEY=
diff --git a/pages/api/generate.ts b/pages/api/generate.ts
index 6f3cad69..dc51b67a 100644
--- a/pages/api/generate.ts
+++ b/pages/api/generate.ts
@@ -3,7 +3,12 @@ import type { NextApiRequest, NextApiResponse } from "next";
 import requestIp from "request-ip";
 import redis from "../../utils/redis";
 
-type Data = string;
+export type GenerateResponseData = {
+  original: string | null;
+  generated: string | null;
+  id: string;
+};
+
 interface ExtendedNextApiRequest extends NextApiRequest {
   body: {
     imageUrl: string;
@@ -23,7 +28,7 @@ const ratelimit = redis
 
 export default async function handler(
   req: ExtendedNextApiRequest,
-  res: NextApiResponse<Data>
+  res: NextApiResponse<GenerateResponseData | string>
 ) {
   // Rate Limiter Code
   if (ratelimit) {
@@ -68,10 +73,12 @@ export default async function handler(
   let jsonStartResponse = await startResponse.json();
 
   let endpointUrl = jsonStartResponse.urls.get;
+  const originalImage = jsonStartResponse.input.image;
+  const roomId = jsonStartResponse.id;
 
   // GET request to get the status of the image restoration process & return the result when it's ready
-  let restoredImage: string | null = null;
-  while (!restoredImage) {
+  let generatedImage: string | null = null;
+  while (!generatedImage) {
     // Loop in 1s intervals until the alt text is ready
     console.log("polling for result...");
     let finalResponse = await fetch(endpointUrl, {
@@ -84,14 +91,20 @@ export default async function handler(
     let jsonFinalResponse = await finalResponse.json();
 
     if (jsonFinalResponse.status === "succeeded") {
-      restoredImage = jsonFinalResponse.output;
+      generatedImage = jsonFinalResponse.output[1] as string;
     } else if (jsonFinalResponse.status === "failed") {
       break;
     } else {
       await new Promise((resolve) => setTimeout(resolve, 1000));
     }
   }
-  res
-    .status(200)
-    .json(restoredImage ? restoredImage : "Failed to restore image");
+  res.status(200).json(
+    generatedImage
+      ? {
+          original: originalImage,
+          generated: generatedImage,
+          id: roomId,
+        }
+      : "Failed to restore image"
+  );
 }
diff --git a/pages/dream.tsx b/pages/dream.tsx
index 1a36fb94..048501e5 100644
--- a/pages/dream.tsx
+++ b/pages/dream.tsx
@@ -16,6 +16,7 @@ import downloadPhoto from "../utils/downloadPhoto";
 import DropDown from "../components/DropDown";
 import { roomType, rooms, themeType, themes } from "../utils/dropdownTypes";
 import CountUp from "react-countup";
+import { GenerateResponseData } from "./api/generate";
 
 // Configuration for the uploader
 const uploader = Uploader({
@@ -83,11 +84,16 @@ const Home: NextPage = () => {
       body: JSON.stringify({ imageUrl: fileUrl, theme, room }),
     });
 
-    let newPhoto = await res.json();
+    let response = (await res.json()) as GenerateResponseData;
+    console.log("res from replicate", response);
     if (res.status !== 200) {
-      setError(newPhoto);
+      setError(response as any);
     } else {
-      setRestoredImage(newPhoto[1]);
+      const rooms =
+        (JSON.parse(localStorage.getItem("rooms") || "[]") as string[]) || [];
+      rooms.push(response.id);
+      localStorage.setItem("rooms", JSON.stringify(rooms));
+      setRestoredImage(response.generated);
     }
     setTimeout(() => {
       setLoading(false);