From 99e3097122131fed1a4723006964971623351f94 Mon Sep 17 00:00:00 2001
From: Aaryan Khandelwal <65252264+aaryan610@users.noreply.github.com>
Date: Fri, 24 May 2024 20:12:40 +0530
Subject: [PATCH] fix: update theme post sign-in (#4586)

---
 web/helpers/theme.helper.ts        |  3 +--
 web/lib/wrappers/store-wrapper.tsx | 37 +++++++++++++++++++++++-------
 2 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/web/helpers/theme.helper.ts b/web/helpers/theme.helper.ts
index 6ebebc820..fd3bd07af 100644
--- a/web/helpers/theme.helper.ts
+++ b/web/helpers/theme.helper.ts
@@ -59,9 +59,8 @@ const calculateShades = (hexValue: string): TShades => {
   return shades as TShades;
 };
 
-export const applyTheme = (palette: string, isDarkPalette: boolean) => {
+export const applyTheme = (palette: string, isDarkPalette: boolean, dom: HTMLElement | null) => {
   if (!palette) return;
-  const dom = document?.querySelector<HTMLElement>("[data-theme='custom']");
   // palette: [bg, text, primary, sidebarBg, sidebarText]
   const values: string[] = palette.split(",");
   values.push(isDarkPalette ? "dark" : "light");
diff --git a/web/lib/wrappers/store-wrapper.tsx b/web/lib/wrappers/store-wrapper.tsx
index ec5b79bfa..e4e191d5a 100644
--- a/web/lib/wrappers/store-wrapper.tsx
+++ b/web/lib/wrappers/store-wrapper.tsx
@@ -22,7 +22,7 @@ const StoreWrapper: FC<TStoreWrapper> = observer((props) => {
   const { sidebarCollapsed, toggleSidebar } = useAppTheme();
   const { data: userProfile } = useUserProfile();
   // states
-  const [dom, setDom] = useState<undefined | HTMLElement>();
+  const [dom, setDom] = useState<HTMLElement | null>(null);
 
   /**
    * Sidebar collapsed fetching from local storage
@@ -38,19 +38,40 @@ const StoreWrapper: FC<TStoreWrapper> = observer((props) => {
    * Setting up the theme of the user by fetching it from local storage
    */
   useEffect(() => {
-    if (!userProfile?.theme?.theme) return;
-    if (window) setDom(() => window.document?.querySelector<HTMLElement>("[data-theme='custom']") || undefined);
-
     setTheme(userProfile?.theme?.theme || "system");
-    if (userProfile?.theme?.theme === "custom" && userProfile?.theme?.palette && dom)
+    if (!userProfile?.theme?.theme) return;
+
+    if (userProfile?.theme?.theme === "custom" && userProfile?.theme?.palette) {
       applyTheme(
         userProfile?.theme?.palette !== ",,,,"
           ? userProfile?.theme?.palette
           : "#0d101b,#c5c5c5,#3f76ff,#0d101b,#c5c5c5",
-        false
+        false,
+        dom
       );
-    else unsetCustomCssVariables();
-  }, [userProfile, userProfile?.theme?.theme, userProfile?.theme?.palette, setTheme, dom]);
+    } else unsetCustomCssVariables();
+  }, [userProfile, userProfile?.theme, userProfile?.theme?.palette, setTheme, dom]);
+
+  useEffect(() => {
+    if (dom) return;
+
+    const observer = new MutationObserver((mutationsList, observer) => {
+      for (const mutation of mutationsList) {
+        if (mutation.type === "childList") {
+          const customThemeElement = window.document?.querySelector<HTMLElement>("[data-theme='custom']");
+          if (customThemeElement) {
+            setDom(customThemeElement);
+            observer.disconnect();
+            break;
+          }
+        }
+      }
+    });
+
+    observer.observe(document.body, { childList: true, subtree: true });
+
+    return () => observer.disconnect();
+  }, [dom]);
 
   useEffect(() => {
     if (!router.query) return;