superwork/vscode/webview-ui/src/codeReader/useMermaidDiagrams.ts

import type { Mermaid } from 'mermaid'
import { useEffect, useState } from 'react'
import { vscodeThemeKind } from '@/lib/vscodeTheme'

/** mermaid 的配色：VS Code 浅色主题用 `default`，深色和高对比用 `dark`。 */
type DiagramTheme = 'default' | 'dark'

/** 一张图画出来的样子：SVG，或者画不出来的原因。 */
type Drawn = { svg: string } | { error: string }

/** 还没按当前配色画好的一张图。 */
interface Pending {
  source: string
  theme: DiagramTheme
}

/** 画好的图最多记几张，最久没用的先丢。 */
const CACHE_SIZE = 100

function currentTheme(): DiagramTheme {
  return vscodeThemeKind() === 'light' ? 'default' : 'dark'
}

/** 页面正文的字体（VS Code 设置里的界面字体），让图里的字和注释一样。没设就用 mermaid 自己的。 */
function pageFont(): string | undefined {
  return getComputedStyle(document.body).getPropertyValue('--vscode-font-family').trim() || undefined
}

/**
 * 页面上 ```mermaid 代码块的占位（见 [Markdown 排版](vscode/webview-ui/src/codeReader/markdown.ts)）画成图。
 *
 * 按源码 + 配色记下画好的图，卡片重新挂上、滚回来直接用。一次只画一张；画的时候页面或主题又变了，
 * 画完这张按页面现在的样子接着画。第一次遇到占位才加载 mermaid。
 */
class MermaidDiagrams {
  private readonly drawn = new Map<string, Drawn>()
  /** 占位里原来的代码块：换成图以后画不出来了，放回去 */
  private readonly codeBlocks = new WeakMap<HTMLElement, Node[]>()
  private mermaid?: Promise<Mermaid>
  /** mermaid 按哪种配色初始化过 */
  private initialized?: DiagramTheme
  private drawing = false
  private count = 0

  /** 把页面上还没按当前配色画好的占位都画上。正在画就不另起一轮，正在跑的那一轮画完会接着看页面。 */
  async draw(): Promise<void> {
    if (this.drawing)
      return
    this.drawing = true
    try {
      for (let next = this.paint(); next; next = this.paint())
        this.remember(next, await this.render(next))
    }
    finally {
      this.drawing = false
    }
  }

  /** 画过的图直接换上；返回第一张还没画过的，都画过了返回 `undefined`。 */
  private paint(): Pending | undefined {
    const theme = currentTheme()
    let pending: Pending | undefined
    for (const node of document.querySelectorAll<HTMLElement>('.cr-markdown .cr-mermaid')) {
      const source = node.dataset.source
      if (source === undefined || node.dataset.drawn === theme)
        continue
      const drawn = this.recall({ source, theme })
      if (drawn)
        this.show(node, drawn, theme)
      else
        pending ??= { source, theme }
    }
    return pending
  }

  /** 画好了把代码块换成图；画不出来留着代码块，下面一行写原因的第一行，悬停看全文。 */
  private show(node: HTMLElement, drawn: Drawn, theme: DiagramTheme): void {
    const codeBlock = this.codeBlocks.get(node) ?? [...node.childNodes]
    this.codeBlocks.set(node, codeBlock)
    if ('svg' in drawn) {
      node.innerHTML = drawn.svg
    }
    else {
      const line = document.createElement('p')
      line.className = 'cr-mermaid-error'
      line.textContent = `图画不出来：${drawn.error.split('\n')[0]}`
      if (drawn.error.includes('\n'))
        line.dataset.tooltip = drawn.error
      node.replaceChildren(...codeBlock, line)
    }
    node.dataset.drawn = theme
  }

  private async render({ source, theme }: Pending): Promise<Drawn> {
    try {
      const mermaid = await (this.mermaid ??= import('mermaid').then(module => module.default))
      if (this.initialized !== theme) {
        // strict：图里的文字不当 HTML、点击不跑脚本；出错不往页面上画 mermaid 自己的报错图，由 [show] 写原因
        mermaid.initialize({ startOnLoad: false, securityLevel: 'strict', suppressErrorRendering: true, theme, fontFamily: pageFont() })
        this.initialized = theme
      }
      this.count += 1
      const { svg } = await mermaid.render(`cr-mermaid-${this.count}`, source)
      return { svg }
    }
    catch (error) {
      return { error: error instanceof Error ? error.message : String(error) }
    }
  }

  private recall(diagram: Pending): Drawn | undefined {
    const key = `${diagram.theme}\u0000${diagram.source}`
    const drawn = this.drawn.get(key)
    if (drawn) {
      this.drawn.delete(key)
      this.drawn.set(key, drawn)
    }
    return drawn
  }

  private remember(diagram: Pending, drawn: Drawn): void {
    this.drawn.set(`${diagram.theme}\u0000${diagram.source}`, drawn)
    if (this.drawn.size > CACHE_SIZE)
      this.drawn.delete(this.drawn.keys().next().value!)
  }
}

/**
 * 注释、说明、问 AI 的回答里的 ```mermaid 代码块画成图：页面内容变了看有没有新的占位，VS Code 换了主题按新配色重画。
 * 页面上没有占位时只多一次查找，不加载 mermaid。
 */
export function useMermaidDiagrams(): void {
  const [diagrams] = useState(() => new MermaidDiagrams())

  useEffect(() => {
    let timer = 0
    const scan = (): void => {
      window.clearTimeout(timer)
      timer = window.setTimeout(() => void diagrams.draw(), 50)
    }
    const content = new MutationObserver(scan)
    content.observe(document.body, { childList: true, subtree: true })
    const theme = new MutationObserver(scan)
    theme.observe(document.body, { attributes: true, attributeFilter: ['class'] })
    scan()
    return () => {
      content.disconnect()
      theme.disconnect()
      window.clearTimeout(timer)
    }
  }, [diagrams])
}