Update useTheme.mjs

Co-Authored-By: Aviv Keller <me@aviv.sh>

Author: AugustinMauroy

src/generators/web/ui/hooks/useTheme.mjs

@@ -1,112 +1,56 @@
 import { useState, useEffect, useCallback } from 'react';
 
-const THEME_STORAGE_KEY = 'theme';
-const THEME_PREFERENCES = new Set(['system', 'light', 'dark']);
-
-/**
- * Sets up theme toggle button and system preference listener
- */
+/** @returns {'dark'|'light'} The current OS-level color scheme. */
 const getSystemTheme = () =>
   matchMedia('(prefers-color-scheme: dark)').matches ? 'dark' : 'light';
 
 /**
- * Retrieves the stored theme preference from local storage.
- *
- * @returns {'system'|'light'|'dark'|null} The stored theme preference or null if not found.
+ * Applies a theme to the document root.
+ * Resolves 'system' to the actual OS preference before applying.
+ * @param {'system'|'light'|'dark'} pref - The theme preference.
  */
-const getStoredThemePreference = () => {
-  try {
-    const storedTheme = localStorage.getItem(THEME_STORAGE_KEY);
-    return THEME_PREFERENCES.has(storedTheme) ? storedTheme : null;
-  } catch {
-    return null;
-  }
+const applyTheme = pref => {
+  const theme = pref === 'system' ? getSystemTheme() : pref;
+  document.documentElement.setAttribute('data-theme', theme);
+  document.documentElement.style.colorScheme = theme;
 };
 
 /**
- * Stores the theme preference in local storage.
- * If storage is unavailable, it fails silently, allowing the application to continue functioning with an in-memory preference.
- *
- * @param {'system'|'light'|'dark'} themePreference - The theme preference to store.
+ * Applies the system theme to the document root.
  */
-const setStoredThemePreference = themePreference => {
-  try {
-    localStorage.setItem(THEME_STORAGE_KEY, themePreference);
-  } catch {
-    // Ignore storage failures and keep non-persistent in-memory preference.
-  }
-};
-
-/**
- * Applies a theme preference to the document.
- *
- * The persisted preference can be 'system', but the applied document theme is
- * always resolved to either 'light' or 'dark'.
- *
- * @param {'system'|'light'|'dark'} themePreference - Theme preference.
- */
-const applyThemePreference = themePreference => {
-  const resolvedTheme =
-    themePreference === 'system' ? getSystemTheme() : themePreference;
-
-  document.documentElement.setAttribute('data-theme', resolvedTheme);
-  document.documentElement.style.colorScheme = resolvedTheme;
-};
+const applySystemTheme = () => applyTheme('system');
 
 /**
- * A React hook for managing the application's theme preference.
+ * React hook for managing theme preference.
+ * Persists the choice to localStorage and listens for OS theme changes
+ * when set to 'system'.
+ * @returns {['system'|'light'|'dark', (next: 'system'|'light'|'dark') => void]}
  */
 export const useTheme = () => {
-  const [themePreference, setThemePreferenceState] = useState('system');
+  // Read stored preference once on mount; default to 'system'.
+  const [pref, setPref] = useState(() => {
+    return localStorage.getItem('theme') || 'system';
+  });
 
+  // Apply theme on every preference change, and if 'system',
+  // also listen for OS-level color scheme changes.
   useEffect(() => {
-    // Use persisted preference if available, otherwise default to system.
-    const initialPreference = getStoredThemePreference() || 'system';
+    applyTheme(pref);
 
-    applyThemePreference(initialPreference);
-    setThemePreferenceState(initialPreference);
-  }, []);
-
-  /**
-   * Keep the resolved document theme in sync with system changes
-   * whenever the preference is set to 'system'.
-   */
-  useEffect(() => {
-    if (themePreference !== 'system') {
+    if (pref !== 'system') {
       return;
     }
 
-    const mediaQueryList = matchMedia('(prefers-color-scheme: dark)');
-    /**
-     *
-     */
-    const handleSystemThemeChange = () => applyThemePreference('system');
-
-    if ('addEventListener' in mediaQueryList) {
-      mediaQueryList.addEventListener('change', handleSystemThemeChange);
-      return () => {
-        mediaQueryList.removeEventListener('change', handleSystemThemeChange);
-      };
-    }
-
-    mediaQueryList.addListener(handleSystemThemeChange);
-    return () => {
-      mediaQueryList.removeListener(handleSystemThemeChange);
-    };
-  }, [themePreference]);
-
-  /**
-   * Updates the theme preference and applies it immediately.
-   */
-  const setThemePreference = useCallback(nextPreference => {
-    if (!THEME_PREFERENCES.has(nextPreference)) {
-      return;
-    }
+    const mql = matchMedia('(prefers-color-scheme: dark)');
+    mql.addEventListener('change', applySystemTheme);
+    return () => mql.removeEventListener('change', applySystemTheme);
+  }, [pref]);
 
-    setThemePreferenceState(nextPreference);
-    setStoredThemePreference(nextPreference);
-    applyThemePreference(nextPreference);
+  /** Updates the preference in both React state and localStorage. */
+  const setTheme = useCallback(next => {
+    setPref(next);
+    localStorage.setItem('theme', next);
   }, []);
 
-  return [themePreference, setThemePreference];
+  return [pref, setTheme];
 };