ramxcodes · KartikXx · May 16, 2026
diff --git a/src/app/layout.tsx b/src/app/layout.tsx
@@ -9,7 +9,6 @@ import { Poppins, Lexend, Montserrat, Roboto, Inter } from "next/font/google";
 import { ReactLenis } from "@/utils/lenis";
 import { Toaster } from "@/components/ui/sonner";
 import MainNav from "@/components/core/MainNav";
-import Head from "next/head";
 import { ThemeInitScript } from "@/components/ui/ThemeInitScript";
 import { GoogleAnalytics } from "@next/third-parties/google";
 
@@ -153,12 +152,10 @@ export default function RootLayout({
 }>) {
   return (
     <ViewTransitions>
-      <html lang="en">
+      <html lang="en" suppressHydrationWarning>
         {process.env.NEXT_PUBLIC_ENABLE_UMAMI === "true" && <Umami />}
         <ReactLenis root>
-          <Head>
-            <ThemeInitScript />
-          </Head>
+          <ThemeInitScript defaultTheme="light" storageKey="notes-buddy-theme" />
           <body
             className={`${poppins.variable} ${lexend.variable} ${montserrat.variable} ${roboto.variable} ${inter.variable}`}
           >

diff --git a/src/components/ui/ThemeInitScript.tsx b/src/components/ui/ThemeInitScript.tsx
@@ -1,15 +1,43 @@
-export function ThemeInitScript() {
+import Script from "next/script";
+
+type ThemeInitScriptProps = {
+  defaultTheme?: "dark" | "light" | "system";
+  storageKey?: string;
+};
+
+export function ThemeInitScript({
+  defaultTheme = "system",
+  storageKey = "notes-buddy-theme",
+}: ThemeInitScriptProps) {
   const script = `
     (function() {
       try {
-        const storedTheme = localStorage.getItem('theme');
-        if (storedTheme === 'dark' ) {
-          document.documentElement.classList.add('dark');
-        } else {
-          document.documentElement.classList.remove('dark');
+        var storageKey = ${JSON.stringify(storageKey)};
+        var defaultTheme = ${JSON.stringify(defaultTheme)};
+        var storedTheme = localStorage.getItem(storageKey) || defaultTheme;
+        var validTheme = storedTheme === 'dark' || storedTheme === 'light' || storedTheme === 'system';
+
+        if (!validTheme) {
+          storedTheme = defaultTheme;
         }
+
+        var prefersDark = window.matchMedia('(prefers-color-scheme: dark)').matches;
+        var resolvedTheme = storedTheme === 'system'
+          ? (prefersDark ? 'dark' : 'light')
+          : storedTheme;
+
+        document.documentElement.classList.remove('light', 'dark');
+        document.documentElement.classList.add(resolvedTheme);
+        document.documentElement.style.colorScheme = resolvedTheme;
       } catch (e) {}
     })();
   `;
-  return <script dangerouslySetInnerHTML={{ __html: script }} />;
+
+  return (
+    <Script
+      id="theme-init"
+      strategy="beforeInteractive"
+      dangerouslySetInnerHTML={{ __html: script }}
+    />
+  );
 }
diff --git a/src/components/ui/theme-provider.tsx b/src/components/ui/theme-provider.tsx
@@ -23,31 +23,39 @@ const initialState: ThemeProviderState = {
 const ThemeProviderContext =
   React.createContext<ThemeProviderState>(initialState);
 
+const useIsomorphicLayoutEffect =
+  typeof window === "undefined" ? React.useEffect : React.useLayoutEffect;
+
+function isTheme(value: string | null): value is Theme {
+  return value === "dark" || value === "light" || value === "system";
+}
+
 export function ThemeProvider({
   children,
   defaultTheme = "system",
   storageKey = "notes-buddy-theme",
   ...props
 }: ThemeProviderProps) {
-  const [theme, setTheme] = React.useState<Theme>(defaultTheme);
+  const [theme, setTheme] = React.useState<Theme>(() => {
+    if (typeof window === "undefined") {
+      return defaultTheme;
+    }
+
+    try {
+      const savedTheme = localStorage.getItem(storageKey);
+      return isTheme(savedTheme) ? savedTheme : defaultTheme;
+    } catch {
+      return defaultTheme;
+    }
+  });
   const [mounted, setMounted] = React.useState(false);
 
   // Set mounted to true after hydration
   React.useEffect(() => {
     setMounted(true);
   }, []);
 
-  // Load theme from localStorage after mount
-  React.useEffect(() => {
-    if (mounted && typeof window !== "undefined") {
-      const savedTheme = localStorage.getItem(storageKey) as Theme;
-      if (savedTheme) {
-        setTheme(savedTheme);
-      }
-    }
-  }, [mounted, storageKey]);
-
-  React.useEffect(() => {
+  useIsomorphicLayoutEffect(() => {
     if (!mounted) return;
 
     const root = window.document.documentElement;
@@ -61,17 +69,23 @@ export function ThemeProvider({
         : "light";
 
       root.classList.add(systemTheme);
+      root.style.colorScheme = systemTheme;
       return;
     }
 
     root.classList.add(theme);
+    root.style.colorScheme = theme;
   }, [theme, mounted]);
 
   const value = {
     theme,
     setTheme: (theme: Theme) => {
       if (typeof window !== "undefined") {
-        localStorage.setItem(storageKey, theme);
+        try {
+          localStorage.setItem(storageKey, theme);
+        } catch {
+          // Ignore storage failures so the in-memory theme still updates.
+        }
       }
       setTheme(theme);
     },

diff --git a/src/components/ui/theme-toggle.tsx b/src/components/ui/theme-toggle.tsx
@@ -13,17 +13,18 @@ export function ThemeToggle({ className }: ThemeToggleProps) {
   const isDark = theme === "dark";
 
   return (
-    <div
+    <button
+      type="button"
+      aria-label={isDark ? "Switch to light theme" : "Switch to dark theme"}
+      aria-pressed={isDark}
       className={cn(
-        "flex h-8 w-16 cursor-pointer rounded-full p-1 transition-all duration-300",
+        "flex h-8 w-16 cursor-pointer rounded-full p-1 transition-all duration-300 focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:outline-none",
         isDark
           ? "border border-zinc-800 bg-zinc-950"
           : "border border-zinc-200 bg-white",
         className,
       )}
       onClick={() => setTheme(isDark ? "light" : "dark")}
-      role="button"
-      tabIndex={0}
     >
       <div className="flex w-full items-center justify-between">
         <div
@@ -53,6 +54,6 @@ export function ThemeToggle({ className }: ThemeToggleProps) {
           )}
         </div>
       </div>
-    </div>
+    </button>
   );
 }