dify/web/app/components/base/mermaid/index.tsx

import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
import mermaid, { type MermaidConfig } from 'mermaid'
import { useTranslation } from 'react-i18next'
import { ExclamationTriangleIcon } from '@heroicons/react/24/outline'
import { MoonIcon, SunIcon } from '@heroicons/react/24/solid'
import {
  cleanUpSvgCode,
  isMermaidCodeComplete,
  prepareMermaidCode,
  processSvgForTheme,
  svgToBase64,
  waitForDOMElement,
} from './utils'
import LoadingAnim from '@/app/components/base/chat/chat/loading-anim'
import cn from '@/utils/classnames'
import ImagePreview from '@/app/components/base/image-uploader/image-preview'
import { Theme } from '@/types/app'

// Global flags and cache for mermaid
let isMermaidInitialized = false
const diagramCache = new Map<string, string>()
let mermaidAPI: any = null

if (typeof window !== 'undefined')
  mermaidAPI = mermaid.mermaidAPI

// Theme configurations
const THEMES = {
  light: {
    name: 'Light Theme',
    background: '#ffffff',
    primaryColor: '#ffffff',
    primaryBorderColor: '#000000',
    primaryTextColor: '#000000',
    secondaryColor: '#ffffff',
    tertiaryColor: '#ffffff',
    nodeColors: [
      { bg: '#f0f9ff', color: '#0369a1' },
      { bg: '#f0fdf4', color: '#166534' },
      { bg: '#fef2f2', color: '#b91c1c' },
      { bg: '#faf5ff', color: '#7e22ce' },
      { bg: '#fffbeb', color: '#b45309' },
    ],
    connectionColor: '#74a0e0',
  },
  dark: {
    name: 'Dark Theme',
    background: '#1e293b',
    primaryColor: '#334155',
    primaryBorderColor: '#94a3b8',
    primaryTextColor: '#e2e8f0',
    secondaryColor: '#475569',
    tertiaryColor: '#334155',
    nodeColors: [
      { bg: '#164e63', color: '#e0f2fe' },
      { bg: '#14532d', color: '#dcfce7' },
      { bg: '#7f1d1d', color: '#fee2e2' },
      { bg: '#581c87', color: '#f3e8ff' },
      { bg: '#78350f', color: '#fef3c7' },
    ],
    connectionColor: '#60a5fa',
  },
}

/**
 * Initializes mermaid library with default configuration
 */
const initMermaid = () => {
  if (typeof window !== 'undefined' && !isMermaidInitialized) {
    try {
      const config: MermaidConfig = {
        startOnLoad: false,
        fontFamily: 'sans-serif',
        securityLevel: 'loose',
        flowchart: {
          htmlLabels: true,
          useMaxWidth: true,
          curve: 'basis',
          nodeSpacing: 50,
          rankSpacing: 70,
        },
        gantt: {
          titleTopMargin: 25,
          barHeight: 20,
          barGap: 4,
          topPadding: 50,
          leftPadding: 75,
          gridLineStartPadding: 35,
          fontSize: 11,
          numberSectionStyles: 4,
          axisFormat: '%Y-%m-%d',
        },
        mindmap: {
          useMaxWidth: true,
          padding: 10,
        },
        maxTextSize: 50000,
      }
      mermaid.initialize(config)
      isMermaidInitialized = true
    }
    catch (error) {
      console.error('Mermaid initialization error:', error)
      return null
    }
  }
  return isMermaidInitialized
}

const Flowchart = React.forwardRef((props: {
  PrimitiveCode: string
  theme?: 'light' | 'dark'
}, ref) => {
  const { t } = useTranslation()
  const [svgString, setSvgString] = useState<string | null>(null)
  const [look, setLook] = useState<'classic' | 'handDrawn'>('classic')
  const [isInitialized, setIsInitialized] = useState(false)
  const [currentTheme, setCurrentTheme] = useState<'light' | 'dark'>(props.theme || 'light')
  const containerRef = useRef<HTMLDivElement>(null)
  const chartId = useRef(`mermaid-chart-${Math.random().toString(36).substr(2, 9)}`).current
  const [isLoading, setIsLoading] = useState(true)
  const renderTimeoutRef = useRef<NodeJS.Timeout>()
  const [errMsg, setErrMsg] = useState('')
  const [imagePreviewUrl, setImagePreviewUrl] = useState('')
  const [isCodeComplete, setIsCodeComplete] = useState(false)
  const codeCompletionCheckRef = useRef<NodeJS.Timeout>()
  const prevCodeRef = useRef<string>()

  // Create cache key from code, style and theme
  const cacheKey = useMemo(() => {
    return `${props.PrimitiveCode}-${look}-${currentTheme}`
  }, [props.PrimitiveCode, look, currentTheme])

  /**
   * Renders Mermaid chart
   */
  const renderMermaidChart = async (code: string, style: 'classic' | 'handDrawn') => {
    if (style === 'handDrawn') {
      // Special handling for hand-drawn style
      if (containerRef.current)
        containerRef.current.innerHTML = `<div id="${chartId}"></div>`
      await new Promise(resolve => setTimeout(resolve, 30))

      if (typeof window !== 'undefined' && mermaidAPI) {
        // Prefer using mermaidAPI directly for hand-drawn style
        return await mermaidAPI.render(chartId, code)
      }
      else {
        // Fall back to standard rendering if mermaidAPI is not available
        const { svg } = await mermaid.render(chartId, code)
        return { svg }
      }
    }
    else {
      // Standard rendering for classic style - using the extracted waitForDOMElement function
      const renderWithRetry = async () => {
        if (containerRef.current)
          containerRef.current.innerHTML = `<div id="${chartId}"></div>`
        await new Promise(resolve => setTimeout(resolve, 30))
        const { svg } = await mermaid.render(chartId, code)
        return { svg }
      }
      return await waitForDOMElement(renderWithRetry)
    }
  }

  /**
   * Handle rendering errors
   */
  const handleRenderError = (error: any) => {
    console.error('Mermaid rendering error:', error)

    // On any render error, assume the mermaid state is corrupted and force a re-initialization.
    try {
      diagramCache.clear() // Clear cache to prevent using potentially corrupted SVGs
      isMermaidInitialized = false // <-- THE FIX: Force re-initialization
      initMermaid() // Re-initialize with the default safe configuration
    }
    catch (reinitError) {
      console.error('Failed to re-initialize Mermaid after error:', reinitError)
    }

    setErrMsg(`Rendering failed: ${(error as Error).message || 'Unknown error. Please check the console.'}`)
    setIsLoading(false)
  }

  // Initialize mermaid
  useEffect(() => {
    const api = initMermaid()
    if (api)
      setIsInitialized(true)
  }, [])

  // Update theme when prop changes, but allow internal override.
  const prevThemeRef = useRef<string>()
  useEffect(() => {
    // Only react if the theme prop from the outside has actually changed.
    if (props.theme && props.theme !== prevThemeRef.current) {
      // When the global theme prop changes, it should act as the source of truth,
      // overriding any local theme selection.
      diagramCache.clear()
      setSvgString(null)
      setCurrentTheme(props.theme)
      // Reset look to classic for a consistent state after a global change.
      setLook('classic')
    }
    // Update the ref to the current prop value for the next render.
    prevThemeRef.current = props.theme
  }, [props.theme])

  const renderFlowchart = useCallback(async (primitiveCode: string) => {
    if (!isInitialized || !containerRef.current) {
      setIsLoading(false)
      setErrMsg(!isInitialized ? 'Mermaid initialization failed' : 'Container element not found')
      return
    }

    // Return cached result if available
    const cacheKey = `${primitiveCode}-${look}-${currentTheme}`
    if (diagramCache.has(cacheKey)) {
      setErrMsg('')
      setSvgString(diagramCache.get(cacheKey) || null)
      setIsLoading(false)
      return
    }

    setIsLoading(true)
    setErrMsg('')

    try {
      let finalCode: string

      const trimmedCode = primitiveCode.trim()
      const isGantt = trimmedCode.startsWith('gantt')
      const isMindMap = trimmedCode.startsWith('mindmap')
      const isSequence = trimmedCode.startsWith('sequenceDiagram')

      if (isGantt || isMindMap || isSequence) {
        if (isGantt) {
          finalCode = trimmedCode
            .split('\n')
            .map((line) => {
              // Gantt charts have specific syntax needs.
              const taskMatch = line.match(/^\s*([^:]+?)\s*:\s*(.*)/)
              if (!taskMatch)
                return line // Not a task line, return as is.

              const taskName = taskMatch[1].trim()
              let paramsStr = taskMatch[2].trim()

              // Rule 1: Correct multiple "after" dependencies ONLY if they exist.
              // This is a common mistake, e.g., "..., after task1, after task2, ..."
              const afterCount = (paramsStr.match(/after /g) || []).length
              if (afterCount > 1)
                paramsStr = paramsStr.replace(/,\s*after\s+/g, ' ')

              // Rule 2: Normalize spacing between parameters for consistency.
              const finalParams = paramsStr.replace(/\s*,\s*/g, ', ').trim()
              return `${taskName} :${finalParams}`
            })
            .join('\n')
        }
        else {
          // For mindmap and sequence charts, which are sensitive to syntax,
          // pass the code through directly.
          finalCode = trimmedCode
        }
      }
      else {
        // Step 1: Clean and prepare Mermaid code using the extracted prepareMermaidCode function
        // This function handles flowcharts appropriately.
        finalCode = prepareMermaidCode(primitiveCode, look)
      }

      // Step 2: Render chart
      const svgGraph = await renderMermaidChart(finalCode, look)

      // Step 3: Apply theme to SVG using the extracted processSvgForTheme function
      const processedSvg = processSvgForTheme(
        svgGraph.svg,
        currentTheme === Theme.dark,
        look === 'handDrawn',
        THEMES,
      )

      // Step 4: Clean up SVG code
      const cleanedSvg = cleanUpSvgCode(processedSvg)

      if (cleanedSvg && typeof cleanedSvg === 'string') {
        diagramCache.set(cacheKey, cleanedSvg)
        setSvgString(cleanedSvg)
      }

      setIsLoading(false)
    }
    catch (error) {
      // Error handling
      handleRenderError(error)
    }
  }, [chartId, isInitialized, look, currentTheme, t])

  const configureMermaid = useCallback((primitiveCode: string) => {
    if (typeof window !== 'undefined' && isInitialized) {
      const themeVars = THEMES[currentTheme]
      const config: any = {
        startOnLoad: false,
        securityLevel: 'loose',
        fontFamily: 'sans-serif',
        maxTextSize: 50000,
        gantt: {
          titleTopMargin: 25,
          barHeight: 20,
          barGap: 4,
          topPadding: 50,
          leftPadding: 75,
          gridLineStartPadding: 35,
          fontSize: 11,
          numberSectionStyles: 4,
          axisFormat: '%Y-%m-%d',
        },
        mindmap: {
          useMaxWidth: true,
          padding: 10,
        },
      }

      const isFlowchart = primitiveCode.trim().startsWith('graph') || primitiveCode.trim().startsWith('flowchart')

      if (look === 'classic') {
        config.theme = currentTheme === 'dark' ? 'dark' : 'neutral'

        if (isFlowchart) {
          config.flowchart = {
            htmlLabels: true,
            useMaxWidth: true,
            nodeSpacing: 60,
            rankSpacing: 80,
            curve: 'linear',
            ranker: 'tight-tree',
          }
        }

        if (currentTheme === 'dark') {
          config.themeVariables = {
            background: themeVars.background,
            primaryColor: themeVars.primaryColor,
            primaryBorderColor: themeVars.primaryBorderColor,
            primaryTextColor: themeVars.primaryTextColor,
            secondaryColor: themeVars.secondaryColor,
            tertiaryColor: themeVars.tertiaryColor,
          }
        }
      }
      else { // look === 'handDrawn'
        config.theme = 'default'
        config.themeCSS = `
          .node rect { fill-opacity: 0.85; }
          .edgePath .path { stroke-width: 1.5px; }
          .label { font-family: 'sans-serif'; }
          .edgeLabel { font-family: 'sans-serif'; }
          .cluster rect { rx: 5px; ry: 5px; }
        `
        config.themeVariables = {
          fontSize: '14px',
          fontFamily: 'sans-serif',
          primaryBorderColor: currentTheme === 'dark' ? THEMES.dark.connectionColor : THEMES.light.connectionColor,
        }

        if (isFlowchart) {
          config.flowchart = {
            htmlLabels: true,
            useMaxWidth: true,
            nodeSpacing: 40,
            rankSpacing: 60,
            curve: 'basis',
          }
        }
      }

      try {
        mermaid.initialize(config)
        return true
      }
      catch (error) {
        console.error('Config error:', error)
        return false
      }
    }
    return false
  }, [currentTheme, isInitialized, look])

  // This is the main rendering effect.
  // It triggers whenever the code, theme, or style changes.
  useEffect(() => {
    if (!isInitialized)
      return

    // Don't render if code is too short
    if (!props.PrimitiveCode || props.PrimitiveCode.length < 10) {
      setIsLoading(false)
      setSvgString(null)
      return
    }

    // Use a timeout to handle streaming code and debounce rendering
    if (renderTimeoutRef.current)
      clearTimeout(renderTimeoutRef.current)

    setIsLoading(true)

    renderTimeoutRef.current = setTimeout(() => {
      // Final validation before rendering
      if (!isMermaidCodeComplete(props.PrimitiveCode)) {
        setIsLoading(false)
        setErrMsg('Diagram code is not complete or invalid.')
        return
      }

      const cacheKey = `${props.PrimitiveCode}-${look}-${currentTheme}`
      if (diagramCache.has(cacheKey)) {
        setErrMsg('')
        setSvgString(diagramCache.get(cacheKey) || null)
        setIsLoading(false)
        return
      }

      if (configureMermaid(props.PrimitiveCode))
        renderFlowchart(props.PrimitiveCode)
    }, 300) // 300ms debounce

    return () => {
      if (renderTimeoutRef.current)
        clearTimeout(renderTimeoutRef.current)
    }
  }, [props.PrimitiveCode, look, currentTheme, isInitialized, configureMermaid, renderFlowchart])

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (containerRef.current)
        containerRef.current.innerHTML = ''
      if (renderTimeoutRef.current)
        clearTimeout(renderTimeoutRef.current)
    }
  }, [])

  const handlePreviewClick = async () => {
    if (svgString) {
      const base64 = await svgToBase64(svgString)
      setImagePreviewUrl(base64)
    }
  }

  const toggleTheme = () => {
    const newTheme = currentTheme === 'light' ? 'dark' : 'light'
    // Ensure a full, clean re-render cycle, consistent with global theme change.
    diagramCache.clear()
    setSvgString(null)
    setCurrentTheme(newTheme)
  }

  // Style classes for theme-dependent elements
  const themeClasses = {
    container: cn('relative', {
      'bg-white': currentTheme === Theme.light,
      'bg-slate-900': currentTheme === Theme.dark,
    }),
    mermaidDiv: cn('mermaid relative h-auto w-full cursor-pointer', {
      'bg-white': currentTheme === Theme.light,
      'bg-slate-900': currentTheme === Theme.dark,
    }),
    errorMessage: cn('px-[26px] py-4', {
      'text-red-500': currentTheme === Theme.light,
      'text-red-400': currentTheme === Theme.dark,
    }),
    errorIcon: cn('h-6 w-6', {
      'text-red-500': currentTheme === Theme.light,
      'text-red-400': currentTheme === Theme.dark,
    }),
    segmented: cn('msh-segmented msh-segmented-sm css-23bs09 css-var-r1', {
      'text-gray-700': currentTheme === Theme.light,
      'text-gray-300': currentTheme === Theme.dark,
    }),
    themeToggle: cn('flex h-10 w-10 items-center justify-center rounded-full shadow-md backdrop-blur-sm transition-all duration-300', {
      'bg-white/80 hover:bg-white hover:shadow-lg text-gray-700 border border-gray-200': currentTheme === Theme.light,
      'bg-slate-800/80 hover:bg-slate-700 hover:shadow-lg text-yellow-300 border border-slate-600': currentTheme === Theme.dark,
    }),
  }

  // Style classes for look options
  const getLookButtonClass = (lookType: 'classic' | 'handDrawn') => {
    return cn(
      'system-sm-medium mb-4 flex h-8 w-[calc((100%-8px)/2)] cursor-pointer items-center justify-center rounded-lg border border-components-option-card-option-border bg-components-option-card-option-bg text-text-secondary',
      look === lookType && 'border-[1.5px] border-components-option-card-option-selected-border bg-components-option-card-option-selected-bg text-text-primary',
      currentTheme === Theme.dark && 'border-slate-600 bg-slate-800 text-slate-300',
      look === lookType && currentTheme === Theme.dark && 'border-blue-500 bg-slate-700 text-white',
    )
  }

  return (
    <div ref={ref as React.RefObject<HTMLDivElement>} className={themeClasses.container}>
      <div className={themeClasses.segmented}>
        <div className="msh-segmented-group">
          <label className="msh-segmented-item m-2 flex w-[200px] items-center space-x-1">
            <div
              key='classic'
              className={getLookButtonClass('classic')}
              onClick={() => {
                if (look !== 'classic') {
                  diagramCache.clear()
                  setSvgString(null)
                  setLook('classic')
                }
              }}
            >
              <div className="msh-segmented-item-label">{t('app.mermaid.classic')}</div>
            </div>
            <div
              key='handDrawn'
              className={getLookButtonClass('handDrawn')}
              onClick={() => {
                if (look !== 'handDrawn') {
                  diagramCache.clear()
                  setSvgString(null)
                  setLook('handDrawn')
                }
              }}
            >
              <div className="msh-segmented-item-label">{t('app.mermaid.handDrawn')}</div>
            </div>
          </label>
        </div>
      </div>

      <div ref={containerRef} style={{ position: 'absolute', visibility: 'hidden', height: 0, overflow: 'hidden' }} />

      {isLoading && !svgString && (
        <div className='px-[26px] py-4'>
          <LoadingAnim type='text'/>
          {!isCodeComplete && (
            <div className="mt-2 text-sm text-gray-500">
              {t('common.wait_for_completion', 'Waiting for diagram code to complete...')}
            </div>
          )}
        </div>
      )}

      {svgString && (
        <div className={themeClasses.mermaidDiv} style={{ objectFit: 'cover' }} onClick={handlePreviewClick}>
          <div className="absolute bottom-2 left-2 z-[100]">
            <button
              onClick={(e) => {
                e.stopPropagation()
                toggleTheme()
              }}
              className={themeClasses.themeToggle}
              title={(currentTheme === Theme.light ? t('app.theme.switchDark') : t('app.theme.switchLight')) || ''}
              style={{ transform: 'translate3d(0, 0, 0)' }}
            >
              {currentTheme === Theme.light ? <MoonIcon className="h-5 w-5" /> : <SunIcon className="h-5 w-5" />}
            </button>
          </div>

          <div
            style={{ maxWidth: '100%' }}
            dangerouslySetInnerHTML={{ __html: svgString }}
          />
        </div>
      )}

      {errMsg && (
        <div className={themeClasses.errorMessage}>
          <div className="flex items-center">
            <ExclamationTriangleIcon className={themeClasses.errorIcon}/>
            <span className="ml-2">{errMsg}</span>
          </div>
        </div>
      )}

      {imagePreviewUrl && (
        <ImagePreview title='mermaid_chart' url={imagePreviewUrl} onCancel={() => setImagePreviewUrl('')} />
      )}
    </div>
  )
})

Flowchart.displayName = 'Flowchart'

export default Flowchart