From aa1de57d96df0a78ab95cb6ad54928705f72dd94 Mon Sep 17 00:00:00 2001
From: Himself65 <himself65@outlook.com>
Date: Sun, 19 Feb 2023 02:38:43 -0600
Subject: [PATCH] feat: read local theme on page load (#1114)

Co-authored-by: zqran <uuxnet@gmail.com>
---
 apps/web/src/providers/ThemeProvider.tsx      | 55 ++++++++++++++-----
 .../styles/utils/localStorageThemeHelper.ts   |  2 +
 2 files changed, 43 insertions(+), 14 deletions(-)

diff --git a/apps/web/src/providers/ThemeProvider.tsx b/apps/web/src/providers/ThemeProvider.tsx
index e4568fe66b..e875e79418 100644
--- a/apps/web/src/providers/ThemeProvider.tsx
+++ b/apps/web/src/providers/ThemeProvider.tsx
@@ -1,5 +1,4 @@
 import {
-  Theme,
   ThemeMode,
   ThemeProviderProps,
   ThemeProviderValue,
@@ -17,7 +16,15 @@ import {
   ThemeProvider as MuiThemeProvider,
 } from '@mui/material/styles';
 import type { PropsWithChildren } from 'react';
-import { createContext, useContext, useEffect, useState } from 'react';
+import {
+  createContext,
+  useCallback,
+  useContext,
+  useEffect,
+  useMemo,
+  useState,
+  useSyncExternalStore,
+} from 'react';
 
 import useCurrentPageMeta from '@/hooks/use-current-page-meta';
 
@@ -35,16 +42,31 @@ export const ThemeProvider = ({
   defaultTheme = 'light',
   children,
 }: PropsWithChildren<ThemeProviderProps>) => {
-  const [theme, setTheme] = useState<Theme>(defaultTheme);
-  const [mode, setMode] = useState<ThemeMode>('auto');
+  const localStorageThemeMode = useSyncExternalStore<ThemeMode>(
+    useCallback(cb => {
+      localStorageThemeHelper.callback.add(cb);
+      return () => {
+        localStorageThemeHelper.callback.delete(cb);
+      };
+    }, []),
+    useCallback(() => localStorageThemeHelper.get() ?? 'light', []),
+    useCallback(() => defaultTheme, [defaultTheme])
+  );
+  const [mode, setMode] = useState<ThemeMode>(defaultTheme);
+  if (localStorageThemeMode !== mode) {
+    setMode(localStorageThemeMode);
+  }
   const { mode: editorMode = 'page' } = useCurrentPageMeta() || {};
   const themeStyle =
-    theme === 'light' ? getLightTheme(editorMode) : getDarkTheme(editorMode);
-  const changeMode = (themeMode: ThemeMode) => {
-    themeMode !== mode && setMode(themeMode);
-    // Remember the theme mode which user selected for next time
-    localStorageThemeHelper.set(themeMode);
-  };
+    mode === 'light' ? getLightTheme(editorMode) : getDarkTheme(editorMode);
+  const changeMode = useCallback(
+    (themeMode: ThemeMode) => {
+      themeMode !== mode && setMode(themeMode);
+      // Remember the theme mode which user selected for next time
+      localStorageThemeHelper.set(themeMode);
+    },
+    [mode]
+  );
 
   // ===================== A temporary solution, just use system theme and not remember the user selected ====================
   useEffect(() => {
@@ -57,9 +79,9 @@ export const ThemeProvider = ({
     });
   }, []);
 
-  useEffect(() => {
-    setTheme(mode === 'auto' ? theme : mode);
-  }, [mode, setTheme, theme]);
+  // useEffect(() => {
+  //   setTheme(mode === 'auto' ? theme : mode);
+  // }, [mode, setTheme, theme]);
   // =====================  ====================
 
   // useEffect(() => {
@@ -93,7 +115,12 @@ export const ThemeProvider = ({
   return (
     // Use MuiThemeProvider is just because some Transitions in Mui components need it
     <MuiThemeProvider theme={muiTheme}>
-      <ThemeContext.Provider value={{ mode, changeMode, theme: themeStyle }}>
+      <ThemeContext.Provider
+        value={useMemo(
+          () => ({ mode, changeMode, theme: themeStyle }),
+          [changeMode, mode, themeStyle]
+        )}
+      >
         <Global
           styles={css`
             :root {
diff --git a/packages/component/src/styles/utils/localStorageThemeHelper.ts b/packages/component/src/styles/utils/localStorageThemeHelper.ts
index 4aee8eca0f..068defc189 100644
--- a/packages/component/src/styles/utils/localStorageThemeHelper.ts
+++ b/packages/component/src/styles/utils/localStorageThemeHelper.ts
@@ -2,11 +2,13 @@ import { ThemeMode } from '../types';
 
 export class LocalStorageThemeHelper {
   name = 'Affine-theme-mode';
+  callback = new Set<() => void>();
   get = (): ThemeMode | null => {
     return localStorage.getItem(this.name) as ThemeMode | null;
   };
   set = (mode: ThemeMode) => {
     localStorage.setItem(this.name, mode);
+    this.callback.forEach(cb => cb());
   };
 }