🤖 fix: improve storybook theme injection

Author: ammar-agent

.storybook/preview.tsx

@@ -1,25 +1,82 @@
 import React from "react";
 import type { Preview } from "@storybook/react-vite";
-import {
-  ThemeProvider,
-  type ThemeMode,
-} from "../src/browser/contexts/ThemeContext";
+import { ThemeProvider, type ThemeMode } from "../src/browser/contexts/ThemeContext";
+import { getStorageChangeEvent } from "@/common/constants/events";
 import { UI_THEME_KEY } from "@/common/constants/storage";
 import "../src/browser/styles/globals.css";
 
-const syncStorybookTheme = (mode?: ThemeMode) => {
-  if (typeof window === "undefined" || !mode) {
+const applyStorybookTheme = (mode: ThemeMode) => {
+  if (typeof window === "undefined") {
     return;
   }
 
   try {
-    window.localStorage.setItem(UI_THEME_KEY, JSON.stringify(mode));
+    const serialized = JSON.stringify(mode);
+    window.localStorage.setItem(UI_THEME_KEY, serialized);
     const root = document.documentElement;
     root.dataset.theme = mode;
     root.style.colorScheme = mode;
+
+    const event = new CustomEvent(getStorageChangeEvent(UI_THEME_KEY), {
+      detail: { key: UI_THEME_KEY, newValue: mode, origin: "storybook" },
+    });
+    window.dispatchEvent(event);
   } catch (error) {
-    console.warn("Failed to sync Storybook theme:", error);
+    console.warn("Failed to write Storybook theme:", error);
+  }
+};
+
+const syncStorybookTheme = (mode?: ThemeMode): ThemeMode => {
+  if (typeof window === "undefined") {
+    return mode ?? "dark";
+  }
+
+  const stored = window.localStorage.getItem(UI_THEME_KEY);
+  const persisted = stored ? (JSON.parse(stored) as ThemeMode) : undefined;
+  const resolved = mode ?? persisted ?? "dark";
+  applyStorybookTheme(resolved);
+  return resolved;
+};
+
+const StorybookThemeToggle: React.FC<{ initialTheme: ThemeMode }> = ({ initialTheme }) => {
+  const [theme, setTheme] = React.useState(initialTheme);
+
+  React.useEffect(() => {
+    setTheme(initialTheme);
+  }, [initialTheme]);
+
+  const handleToggle = () => {
+    const next = theme === "dark" ? "light" : "dark";
+    applyStorybookTheme(next);
+    setTheme(next);
+  };
+
+  if (typeof window === "undefined") {
+    return null;
   }
+
+  return (
+    <button
+      onClick={handleToggle}
+      style={{
+        position: "fixed",
+        bottom: "1rem",
+        right: "1rem",
+        zIndex: 9999,
+        padding: "0.5rem 1rem",
+        background: theme === "dark" ? "#f5f6f8" : "#1e1e1e",
+        color: theme === "dark" ? "#1e1e1e" : "#f5f6f8",
+        border: "1px solid #777",
+        borderRadius: "4px",
+        cursor: "pointer",
+        fontFamily: "system-ui, sans-serif",
+        fontSize: "12px",
+        boxShadow: "0 2px 5px rgba(0,0,0,0.2)",
+      }}
+    >
+      {theme === "dark" ? "☀️ Light" : "🌙 Dark"}
+    </button>
+  );
 };
 
 const preview: Preview = {
@@ -40,10 +97,11 @@ const preview: Preview = {
   decorators: [
     (Story, context) => {
       const mode = context.globals.theme as ThemeMode | undefined;
-      syncStorybookTheme(mode);
+      const resolved = syncStorybookTheme(mode);
       return (
         <ThemeProvider>
           <Story />
+          {!mode && <StorybookThemeToggle initialTheme={resolved} />}
         </ThemeProvider>
       );
     },