From 785785807408cc250301feefbcfa3a83220e1303 Mon Sep 17 00:00:00 2001
From: Biki Kalita <86558912+Biki-dev@users.noreply.github.com>
Date: Mon, 15 Dec 2025 19:10:21 +0530
Subject: [PATCH] feat: add warning dialog for theme and UI style changes
 (#248)

## Summary

  - Auto-saves diagram to localStorage before theme or UI style changes to prevent data loss
  - Extracts inline handler to `handleDrawioUiChange` for cleaner code
  - Renames `toggleDarkMode` to `handleDarkModeChange` for consistency

  ## Problem

  Changing themes (dark/light) or draw.io UI styles (min/sketch) causes the DrawIoEmbed component to remount, losing all unsaved edits without warning.

  ## Solution

  Added `saveDiagramToStorage()` function that exports the current diagram and saves it to localStorage before any theme/UI change. The existing restore mechanism then loads it back after remount.

  ## Related Issues

  Fixes #243
---
 app/page.tsx                 | 33 +++++++++++++++++++--------------
 contexts/diagram-context.tsx | 26 ++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 14 deletions(-)

diff --git a/app/page.tsx b/app/page.tsx
index 760de71..09ce7a1 100644
--- a/app/page.tsx
+++ b/app/page.tsx
@@ -15,8 +15,13 @@ const drawioBaseUrl =
     process.env.NEXT_PUBLIC_DRAWIO_BASE_URL || "https://embed.diagrams.net"
 
 export default function Home() {
-    const { drawioRef, handleDiagramExport, onDrawioLoad, resetDrawioReady } =
-        useDiagram()
+    const {
+        drawioRef,
+        handleDiagramExport,
+        onDrawioLoad,
+        resetDrawioReady,
+        saveDiagramToStorage,
+    } = useDiagram()
     const [isMobile, setIsMobile] = useState(false)
     const [isChatVisible, setIsChatVisible] = useState(true)
     const [drawioUi, setDrawioUi] = useState<"min" | "sketch">("min")
@@ -35,12 +40,10 @@ export default function Home() {
 
         const savedDarkMode = localStorage.getItem("next-ai-draw-io-dark-mode")
         if (savedDarkMode !== null) {
-            // Use saved preference
             const isDark = savedDarkMode === "true"
             setDarkMode(isDark)
             document.documentElement.classList.toggle("dark", isDark)
         } else {
-            // First visit: match browser preference
             const prefersDark = window.matchMedia(
                 "(prefers-color-scheme: dark)",
             ).matches
@@ -58,12 +61,20 @@ export default function Home() {
         setIsLoaded(true)
     }, [])
 
-    const toggleDarkMode = () => {
+    const handleDarkModeChange = async () => {
+        await saveDiagramToStorage()
         const newValue = !darkMode
         setDarkMode(newValue)
         localStorage.setItem("next-ai-draw-io-dark-mode", String(newValue))
         document.documentElement.classList.toggle("dark", newValue)
-        // Reset so onDrawioLoad fires again after remount
+        resetDrawioReady()
+    }
+
+    const handleDrawioUiChange = async () => {
+        await saveDiagramToStorage()
+        const newUi = drawioUi === "min" ? "sketch" : "min"
+        localStorage.setItem("drawio-theme", newUi)
+        setDrawioUi(newUi)
         resetDrawioReady()
     }
 
@@ -182,15 +193,9 @@ export default function Home() {
                             isVisible={isChatVisible}
                             onToggleVisibility={toggleChatPanel}
                             drawioUi={drawioUi}
-                            onToggleDrawioUi={() => {
-                                const newUi =
-                                    drawioUi === "min" ? "sketch" : "min"
-                                localStorage.setItem("drawio-theme", newUi)
-                                setDrawioUi(newUi)
-                                resetDrawioReady()
-                            }}
+                            onToggleDrawioUi={handleDrawioUiChange}
                             darkMode={darkMode}
-                            onToggleDarkMode={toggleDarkMode}
+                            onToggleDarkMode={handleDarkModeChange}
                             isMobile={isMobile}
                             onCloseProtectionChange={setCloseProtection}
                         />
diff --git a/contexts/diagram-context.tsx b/contexts/diagram-context.tsx
index 0f68863..69abdc1 100644
--- a/contexts/diagram-context.tsx
+++ b/contexts/diagram-context.tsx
@@ -23,6 +23,7 @@ interface DiagramContextType {
         format: ExportFormat,
         sessionId?: string,
     ) => void
+    saveDiagramToStorage: () => Promise<void>
     isDrawioReady: boolean
     onDrawioLoad: () => void
     resetDrawioReady: () => void
@@ -82,6 +83,30 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
         }
     }
 
+    // Save current diagram to localStorage (used before theme/UI changes)
+    const saveDiagramToStorage = async (): Promise<void> => {
+        if (!drawioRef.current) return
+
+        try {
+            const currentXml = await Promise.race([
+                new Promise<string>((resolve) => {
+                    resolverRef.current = resolve
+                    drawioRef.current?.exportDiagram({ format: "xmlsvg" })
+                }),
+                new Promise<string>((_, reject) =>
+                    setTimeout(() => reject(new Error("Export timeout")), 2000),
+                ),
+            ])
+
+            // Only save if diagram has meaningful content (not empty template)
+            if (currentXml && currentXml.length > 300) {
+                localStorage.setItem(STORAGE_DIAGRAM_XML_KEY, currentXml)
+            }
+        } catch (error) {
+            console.error("Failed to save diagram to storage:", error)
+        }
+    }
+
     const loadDiagram = (
         chart: string,
         skipValidation?: boolean,
@@ -280,6 +305,7 @@ export function DiagramProvider({ children }: { children: React.ReactNode }) {
                 handleDiagramExport,
                 clearDiagram,
                 saveDiagramToFile,
+                saveDiagramToStorage,
                 isDrawioReady,
                 onDrawioLoad,
                 resetDrawioReady,