From 032ba62cbefc41df328d8a3f1a767e6115736da5 Mon Sep 17 00:00:00 2001
From: "Tom (plebeius.eth)" <tom@plebbit.com>
Date: Sat, 22 Jun 2024 12:36:38 +0200
Subject: [PATCH] fix(themes): changing theme would bug out

---
 src/hooks/use-theme.ts        | 55 ++++++++++++++++++++---------------
 src/stores/use-theme-store.ts |  3 +-
 2 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/src/hooks/use-theme.ts b/src/hooks/use-theme.ts
index 9ad92adc..b06f815d 100644
--- a/src/hooks/use-theme.ts
+++ b/src/hooks/use-theme.ts
@@ -4,44 +4,53 @@ import { isAllView, isSubscriptionsView } from '../lib/utils/view-utils';
 import useThemeStore from '../stores/use-theme-store';
 import useInitialTheme from './use-initial-theme';
 
+const updateThemeClass = (newTheme: string) => {
+  document.body.classList.remove('yotsuba', 'yotsuba-b', 'futaba', 'burichan', 'tomorrow', 'photon');
+  if (newTheme) {
+    document.body.classList.add(newTheme);
+  }
+};
+
 const useTheme = (): [string, (theme: string) => void] => {
   const location = useLocation();
   const params = useParams();
   const setThemeStore = useThemeStore((state) => state.setTheme);
+  const getTheme = useThemeStore((state) => state.getTheme);
   const loadThemes = useThemeStore((state) => state.loadThemes);
 
+  const initialTheme = useInitialTheme();
+  const [theme, setLocalTheme] = useState<string>(() => initialTheme);
   const [themesLoaded, setThemesLoaded] = useState(false);
 
   useEffect(() => {
-    const load = async () => {
+    const loadAndApplyThemes = async () => {
       await loadThemes();
       setThemesLoaded(true);
     };
-    load();
+
+    loadAndApplyThemes();
   }, [loadThemes]);
 
-  const initialTheme = useInitialTheme();
-
-  const [theme, setLocalTheme] = useState(initialTheme);
-
   useEffect(() => {
-    const previousTheme = document.body.className;
-    document.body.classList.add(initialTheme);
-    return () => {
-      document.body.classList.remove(initialTheme);
-      if (previousTheme) {
-        document.body.classList.add(previousTheme);
-      }
-    };
-  }, [initialTheme]);
+    if (!themesLoaded) return;
 
-  useEffect(() => {
-    if (themesLoaded) {
-      document.body.classList.remove('yotsuba', 'yotsuba-b', 'futaba', 'burichan', 'tomorrow', 'photon');
-      document.body.classList.add(initialTheme);
-      setLocalTheme(initialTheme);
+    const subplebbitAddress = params?.subplebbitAddress;
+    const isInAllView = isAllView(location.pathname, params);
+    const isInSubscriptionsView = isSubscriptionsView(location.pathname, params);
+
+    let storedTheme = null;
+    if (subplebbitAddress) {
+      storedTheme = getTheme(subplebbitAddress);
+    } else if (isInAllView) {
+      storedTheme = getTheme('all');
+    } else if (isInSubscriptionsView) {
+      storedTheme = getTheme('subscriptions');
     }
-  }, [initialTheme, themesLoaded]);
+
+    const themeToSet = storedTheme || initialTheme;
+    setLocalTheme(themeToSet);
+    updateThemeClass(themeToSet);
+  }, [initialTheme, location.pathname, params, getTheme, themesLoaded]);
 
   const setSubplebbitTheme = (newTheme: string) => {
     const subplebbitAddress = params?.subplebbitAddress;
@@ -55,9 +64,9 @@ const useTheme = (): [string, (theme: string) => void] => {
     } else if (isInSubscriptionsView) {
       setThemeStore('subscriptions', newTheme);
     }
+
     setLocalTheme(newTheme);
-    document.body.classList.remove('yotsuba', 'yotsuba-b', 'futaba', 'burichan', 'tomorrow', 'photon');
-    document.body.classList.add(newTheme);
+    updateThemeClass(newTheme);
   };
 
   return [theme, setSubplebbitTheme];
diff --git a/src/stores/use-theme-store.ts b/src/stores/use-theme-store.ts
index d20d9609..85c68e91 100644
--- a/src/stores/use-theme-store.ts
+++ b/src/stores/use-theme-store.ts
@@ -5,7 +5,7 @@ interface ThemeState {
   themes: Record<string, string>;
   setTheme: (subplebbitAddress: string, theme: string) => void;
   getTheme: (subplebbitAddress: string) => string | null;
-  loadThemes: () => void;
+  loadThemes: () => Promise<void>;
 }
 
 const themeStore = localForageLru.createInstance({
@@ -35,6 +35,7 @@ const useThemeStore = create<ThemeState>((set: StoreApi<ThemeState>['setState'],
   },
 }));
 
+// Load themes on store initialization
 useThemeStore.getState().loadThemes();
 
 export default useThemeStore;