fix: tweaking theme provider

src/components/theme-provider.tsx

@@ -26,34 +26,42 @@ export function ThemeProvider({
   storageKey = 'vite-ui-theme',
   ...props
 }: ThemeProviderProps) {
-  const [theme, setTheme] = useState<Theme>(
+  const [theme, _setTheme] = useState<Theme>(
     () => (localStorage.getItem(storageKey) as Theme) || defaultTheme
   )
 
   useEffect(() => {
     const root = window.document.documentElement
+    const mediaQuery = window.matchMedia('(prefers-color-scheme: dark)')
 
-    root.classList.remove('light', 'dark')
-
-    if (theme === 'system') {
-      const systemTheme = window.matchMedia('(prefers-color-scheme: dark)')
-        .matches
-        ? 'dark'
-        : 'light'
+    const applyTheme = (theme: Theme) => {
+      root.classList.remove('light', 'dark') // Remove existing theme classes
+      const systemTheme = mediaQuery.matches ? 'dark' : 'light'
+      const effectiveTheme = theme === 'system' ? systemTheme : theme
+      root.classList.add(effectiveTheme) // Add the new theme class
+    }
 
-      root.classList.add(systemTheme)
-      return
+    const handleChange = () => {
+      if (theme === 'system') {
+        applyTheme('system')
+      }
     }
 
-    root.classList.add(theme)
+    applyTheme(theme)
+
+    mediaQuery.addEventListener('change', handleChange)
+
+    return () => mediaQuery.removeEventListener('change', handleChange)
   }, [theme])
 
+  const setTheme = (theme: Theme) => {
+    localStorage.setItem(storageKey, theme)
+    _setTheme(theme)
+  }
+
   const value = {
     theme,
-    setTheme: (theme: Theme) => {
-      localStorage.setItem(storageKey, theme)
-      setTheme(theme)
-    },
+    setTheme,
   }
 
   return (