From 935b7740342daf13dbec6d9564591590870f3783 Mon Sep 17 00:00:00 2001
From: Elroy Chua <elroymx@gmail.com>
Date: Mon, 16 Dec 2024 10:11:59 +0800
Subject: [PATCH] hotfix localstorage state management for schemes and
 userQuery

---
 frontend/src/app/providers.tsx                | 42 ++++++++++++++++---
 .../src/components/main-chat/main-chat.tsx    | 41 +++++++++++-------
 .../src/components/user-query/user-query.tsx  |  4 +-
 3 files changed, 65 insertions(+), 22 deletions(-)

diff --git a/frontend/src/app/providers.tsx b/frontend/src/app/providers.tsx
index 5e6a69f..73d3ddc 100644
--- a/frontend/src/app/providers.tsx
+++ b/frontend/src/app/providers.tsx
@@ -31,17 +31,47 @@ export const ChatProvider = ({ children }: { children: ReactNode }) => {
   const [messages, setMessages] = useState<Message[]>([]);
   const [userQuery, setUserQuery] = useState<string>("");
   const [schemes, setSchemes] = useState<SearchResScheme[]>([]);
+  const [isInitialized, setIsInitialized] = useState(false);
 
   useEffect(() => {
-    const storedSchemes = localStorage.getItem('schemes');
-    if (storedSchemes) {
-      setSchemes(JSON.parse(storedSchemes));
+    if (!isInitialized) {
+      try {
+        const storedSchemes = localStorage.getItem("schemes");
+        const storedQuery = localStorage.getItem("userQuery");
+
+        if (storedSchemes) {
+          const parsedSchemes = JSON.parse(storedSchemes);
+          setSchemes(parsedSchemes);
+        }
+        if (storedQuery) {
+          setUserQuery(storedQuery);
+        }
+        setIsInitialized(true);
+      } catch (error) {
+        console.error("Error loading from localStorage:", error);
+      }
     }
-  }, []);
+  }, [isInitialized]);
 
   useEffect(() => {
-    localStorage.setItem("schemes", JSON.stringify(schemes));
-  }, [schemes]);
+    if (isInitialized) {
+      try {
+        localStorage.setItem("schemes", JSON.stringify(schemes));
+      } catch (error) {
+        console.error("Error saving schemes to localStorage:", error);
+      }
+    }
+  }, [schemes, isInitialized]);
+
+  useEffect(() => {
+    if (isInitialized) {
+      try {
+        localStorage.setItem("userQuery", userQuery);
+      } catch (error) {
+        console.error("Error saving userQuery to localStorage:", error);
+      }
+    }
+  }, [userQuery, isInitialized]);
 
   return (
     <ChatContext.Provider
diff --git a/frontend/src/components/main-chat/main-chat.tsx b/frontend/src/components/main-chat/main-chat.tsx
index 76a10b2..ab0ada8 100644
--- a/frontend/src/components/main-chat/main-chat.tsx
+++ b/frontend/src/components/main-chat/main-chat.tsx
@@ -1,6 +1,6 @@
 "use client";
 
-import { Message, useChat } from "@/app/providers";
+import { useChat } from "@/app/providers";
 import ChatBar from "@/components/chat-bar/chat-bar";
 import ChatList from "@/components/chat-list/chat-list";
 import { Spacer } from "@nextui-org/react";
@@ -13,26 +13,38 @@ type MainChatProps = {
 };
 
 export default function MainChat({ sessionId }: MainChatProps) {
-  const { messages, setMessages } = useChat();
-
+  const { messages, setMessages, setUserQuery } = useChat();
   const [userInput, setUserInput] = useState("");
   const [isBotResponseGenerating, setIsBotResponseGenerating] =
     useState<boolean>(false);
   const [currentStreamingMessage, setCurrentStreamingMessage] = useState("");
-
   const scrollableDivRef = useRef<HTMLDivElement>(null);
 
+  useEffect(() => {
+    const storedQuery = localStorage.getItem("userQuery");
+    if (storedQuery && messages.length === 0) {
+      setMessages([
+        {
+          type: "bot",
+          text: "You can see the search results on the right. Please ask me any further questions about the schemes.",
+        },
+      ]);
+    }
+  }, []);
+
   useEffect(() => {
     handleScrollToBottom();
   }, [messages]);
 
   const handleUserInput = async (input: string) => {
-    setMessages(
-      (prevMessages: Message[]) =>
-        [...prevMessages, { type: "user", text: input }] as Message[]
-    );
+    setMessages((prevMessages) => [
+      ...prevMessages,
+      { type: "user", text: input },
+    ]);
+
+    setUserQuery(input);
+    localStorage.setItem("userQuery", input);
     setUserInput("");
-    // Trigger API call for bot response
     await fetchBotResponse(input);
   };
 
@@ -45,8 +57,7 @@ export default function MainChat({ sessionId }: MainChatProps) {
 
   const fetchBotResponse = async (userMessage: string) => {
     setIsBotResponseGenerating(true);
-    setCurrentStreamingMessage(""); // Reset streaming message
-
+    setCurrentStreamingMessage("");
     try {
       const response = await fetch(
         `${process.env.NEXT_PUBLIC_API_BASE_URL}/chat_message`,
@@ -74,7 +85,7 @@ export default function MainChat({ sessionId }: MainChatProps) {
         throw new Error("No reader available");
       }
 
-      let fullMessage = ""; // Keep track of the full message
+      let fullMessage = "";
 
       while (true) {
         const { done, value } = await reader.read();
@@ -87,8 +98,8 @@ export default function MainChat({ sessionId }: MainChatProps) {
           if (line.startsWith("data: ")) {
             try {
               const data = JSON.parse(line.slice(6));
-              fullMessage += data.chunk; // Add to full message
-              setCurrentStreamingMessage(fullMessage); // Update streaming with full message
+              fullMessage += data.chunk;
+              setCurrentStreamingMessage(fullMessage);
             } catch (e) {
               console.error("Error parsing SSE data:", e);
             }
@@ -101,7 +112,7 @@ export default function MainChat({ sessionId }: MainChatProps) {
       handleBotResponse("Sorry, something went wrong. Please try again.");
     } finally {
       setIsBotResponseGenerating(false);
-      setCurrentStreamingMessage(""); // Clear streaming message after adding to chat history
+      setCurrentStreamingMessage("");
     }
   };
 
diff --git a/frontend/src/components/user-query/user-query.tsx b/frontend/src/components/user-query/user-query.tsx
index 107376b..da1223a 100644
--- a/frontend/src/components/user-query/user-query.tsx
+++ b/frontend/src/components/user-query/user-query.tsx
@@ -6,12 +6,14 @@ import ResetQueryModal from "../reset-query-modal/reset-query-modal";
 import classes from "./user-query.module.css";
 
 export default function UserQuery() {
-  const { userQuery, setSchemes } = useChat();
+  const { userQuery, setSchemes, setUserQuery } = useChat();
   const { isOpen, onOpen, onOpenChange } = useDisclosure();
 
   const handleReset = () => {
     localStorage.removeItem("schemes");
+    localStorage.removeItem("userQuery");
     setSchemes([]);
+    setUserQuery("");
   };
 
   return (