From af1fd901182b4e5c33da8df7b65dc404b2110c55 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Sat, 9 Nov 2024 14:46:35 +0800
Subject: [PATCH] feat: add mermaid library and theme initialization

---
 .../src/context/SyntaxHighlighterProvider.tsx | 28 ++---------
 src/renderer/src/hooks/useMermaid.ts          | 50 +++++++++++++++++++
 src/renderer/src/store/runtime.ts             |  9 ----
 3 files changed, 55 insertions(+), 32 deletions(-)
 create mode 100644 src/renderer/src/hooks/useMermaid.ts

diff --git a/src/renderer/src/context/SyntaxHighlighterProvider.tsx b/src/renderer/src/context/SyntaxHighlighterProvider.tsx
index 1ad7215167..650146f725 100644
--- a/src/renderer/src/context/SyntaxHighlighterProvider.tsx
+++ b/src/renderer/src/context/SyntaxHighlighterProvider.tsx
@@ -1,7 +1,7 @@
 import { useTheme } from '@renderer/context/ThemeProvider'
+import { useMermaid } from '@renderer/hooks/useMermaid'
 import { useSettings } from '@renderer/hooks/useSettings'
 import { CodeStyleVarious, ThemeMode } from '@renderer/types'
-import { loadScript } from '@renderer/utils'
 import React, { createContext, PropsWithChildren, useContext, useEffect, useMemo, useState } from 'react'
 import {
   BundledLanguage,
@@ -22,6 +22,7 @@ export const SyntaxHighlighterProvider: React.FC<PropsWithChildren> = ({ childre
   const { theme } = useTheme()
   const [highlighter, setHighlighter] = useState<HighlighterGeneric<BundledLanguage, BundledTheme> | null>(null)
   const { codeStyle } = useSettings()
+  useMermaid()
 
   const highlighterTheme = useMemo(() => {
     if (!codeStyle || codeStyle === 'auto') {
@@ -31,21 +32,6 @@ export const SyntaxHighlighterProvider: React.FC<PropsWithChildren> = ({ childre
     return codeStyle
   }, [theme, codeStyle])
 
-  useEffect(() => {
-    const initMermaid = async () => {
-      if (!window.mermaid) {
-        await loadScript('https://unpkg.com/mermaid@11.4.0/dist/mermaid.min.js')
-        window.mermaid.initialize({
-          startOnLoad: false,
-          theme: theme === ThemeMode.dark ? 'dark' : 'default'
-        })
-        window.mermaid.contentLoaded()
-      }
-    }
-
-    initMermaid()
-  }, [theme])
-
   useEffect(() => {
     const initHighlighter = async () => {
       const commonLanguages = ['javascript', 'typescript', 'python', 'java', 'markdown']
@@ -57,13 +43,9 @@ export const SyntaxHighlighterProvider: React.FC<PropsWithChildren> = ({ childre
 
       setHighlighter(hl)
 
-      window.requestIdleCallback(
-        () => {
-          hl.loadTheme(...(Object.keys(bundledThemes) as BundledTheme[]))
-          hl.loadLanguage(...(Object.keys(bundledLanguages) as BundledLanguage[]))
-        },
-        { timeout: 2000 }
-      )
+      // Load all themes and languages
+      // hl.loadTheme(...(Object.keys(bundledThemes) as BundledTheme[]))
+      // hl.loadLanguage(...(Object.keys(bundledLanguages) as BundledLanguage[]))
     }
 
     initHighlighter()
diff --git a/src/renderer/src/hooks/useMermaid.ts b/src/renderer/src/hooks/useMermaid.ts
new file mode 100644
index 0000000000..c4ac40dba5
--- /dev/null
+++ b/src/renderer/src/hooks/useMermaid.ts
@@ -0,0 +1,50 @@
+import { useTheme } from '@renderer/context/ThemeProvider'
+import { ThemeMode } from '@renderer/types'
+import { loadScript, runAsyncFunction } from '@renderer/utils'
+import { useEffect } from 'react'
+
+import { useRuntime } from './useRuntime'
+
+export const useMermaid = () => {
+  const { theme } = useTheme()
+  const { generating } = useRuntime()
+
+  useEffect(() => {
+    runAsyncFunction(async () => {
+      if (!window.mermaid) {
+        await loadScript('https://unpkg.com/mermaid@11.4.0/dist/mermaid.min.js')
+        window.mermaid.initialize({
+          startOnLoad: true,
+          theme: theme === ThemeMode.dark ? 'dark' : 'default'
+        })
+        window.mermaid.contentLoaded()
+      }
+    })
+  }, [])
+
+  useEffect(() => {
+    if (window.mermaid) {
+      window.mermaid.initialize({
+        startOnLoad: true,
+        theme: theme === ThemeMode.dark ? 'dark' : 'default'
+      })
+      window.mermaid.contentLoaded()
+    }
+  }, [theme])
+
+  useEffect(() => {
+    if (!window.mermaid || generating) return
+
+    const renderMermaid = () => {
+      const mermaidElements = document.querySelectorAll('.mermaid')
+      mermaidElements.forEach((element) => {
+        if (!element.querySelector('svg')) {
+          element.removeAttribute('data-processed')
+        }
+      })
+      window.mermaid.contentLoaded()
+    }
+
+    setTimeout(renderMermaid, 100)
+  }, [generating])
+}
diff --git a/src/renderer/src/store/runtime.ts b/src/renderer/src/store/runtime.ts
index df7a19f0a0..06eb72d3c2 100644
--- a/src/renderer/src/store/runtime.ts
+++ b/src/renderer/src/store/runtime.ts
@@ -26,15 +26,6 @@ const runtimeSlice = createSlice({
     },
     setGenerating: (state, action: PayloadAction<boolean>) => {
       state.generating = action.payload
-      if (!state.generating) {
-        const mermaidElements = document.querySelectorAll('.mermaid')
-        for (const element of mermaidElements) {
-          if (!element.querySelector('svg')) {
-            element.removeAttribute('data-processed')
-          }
-        }
-        setTimeout(() => window.mermaid.contentLoaded(), 100)
-      }
     },
     setMinappShow: (state, action: PayloadAction<boolean>) => {
       state.minappShow = action.payload