llm-in-text/src/plugins/mermaidPlugin.ts

import { codeBlockConfig } from '@milkdown/kit/component/code-block'
import mermaid from 'mermaid'

let mermaidReadyTheme = ''
let diagramCounter = 0
let renderCounter = 0

type MermaidImagePayload = {
    previewUrl: string
    downloadUrl: string | null
    width: number
    height: number
    sourceWidth: number
    sourceHeight: number
    filename: string
}

function getMermaidTheme() {
    const rootTheme = document.documentElement.getAttribute('data-theme')
    return rootTheme === 'dark' ? 'dark' : 'default'
}

function ensureMermaid() {
    const theme = getMermaidTheme()
    if (mermaidReadyTheme === theme) return

    const dark = theme === 'dark'
    mermaid.initialize({
        startOnLoad: false,
        theme: dark ? 'dark' : 'base',
        securityLevel: 'loose',
        fontFamily: 'inherit',
        flowchart: { htmlLabels: false },
        themeVariables: dark ? {
            primaryColor: '#1e2d45',
            primaryTextColor: '#c9d6e8',
            primaryBorderColor: '#3b5278',
            lineColor: '#5a7aa8',
            secondaryColor: '#162236',
            tertiaryColor: '#0f1926',
            edgeLabelBackground: '#1a2a40',
            clusterBkg: '#111e2e',
            titleColor: '#c9d6e8',
            nodeBorder: '#3b5278',
            mainBkg: '#1e2d45',
        } : {
            primaryColor: '#e8f0fe',
            primaryTextColor: '#1e3a5f',
            primaryBorderColor: '#93b4d9',
            lineColor: '#4a7cb5',
            secondaryColor: '#dbeafe',
            tertiaryColor: '#f0f7ff',
            edgeLabelBackground: '#f0f7ff',
            clusterBkg: '#f5f8ff',
            titleColor: '#1e3a5f',
            nodeBorder: '#93b4d9',
            mainBkg: '#e8f0fe',
        },
    })
    mermaidReadyTheme = theme
}

function encodeMermaidCode(code: string) {
    return encodeURIComponent(code)
}

function decodeMermaidCode(code: string) {
    try {
        return decodeURIComponent(code)
    } catch {
        return code
    }
}

function escapeHtml(value: string) {
    return value
        .replace(/&/g, '&')
        .replace(/</g, '&lt;')
        .replace(/>/g, '&gt;')
        .replace(/"/g, '&quot;')
        .replace(/'/g, '&#39;')
}

function nowFileStamp() {
    const now = new Date()
    const pad = (n: number) => String(n).padStart(2, '0')
    return `${now.getFullYear()}${pad(now.getMonth() + 1)}${pad(now.getDate())}-${pad(now.getHours())}${pad(now.getMinutes())}${pad(now.getSeconds())}`
}

function makeMermaidFilename() {
    return `mermaid-${nowFileStamp()}.png`
}

function buildMermaidPreviewMarkup(code: string, token: number) {
    // 剥离首尾的 ```mermaid 或 ``` 标识符，防止其被误认为图表节点
    const cleanCode = code
        .replace(/^```[a-z]*\s*\n?/i, '')
        .replace(/\n?```\s*$/i, '')
        .trim()
    
    const encoded = encodeMermaidCode(cleanCode)
    const filename = makeMermaidFilename()

    const zoomSvg = `<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.2" stroke-linecap="round" stroke-linejoin="round"><circle cx="11" cy="11" r="8"/><line x1="21" y1="21" x2="16.65" y2="16.65"/><line x1="11" y1="8" x2="11" y2="14"/><line x1="8" y1="11" x2="14" y2="11"/></svg>`
    const dlSvg = `<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2.2" stroke-linecap="round" stroke-linejoin="round"><path d="M21 15v4a2 2 0 0 1-2 2H5a2 2 0 0 1-2-2v-4"/><polyline points="7 10 12 15 17 10"/><line x1="12" y1="15" x2="12" y2="3"/></svg>`

    return `
<div class="mermaid-block" data-mermaid-code="${encoded}" data-mermaid-token="${token}">
  <div class="mermaid-controls">
    <button type="button" class="mermaid-action-btn mermaid-zoom-btn" data-mermaid-action="zoom" disabled aria-label="放大查看">${zoomSvg}</button>
    <button type="button" class="mermaid-action-btn mermaid-download-btn" data-mermaid-action="download" data-mermaid-filename="${filename}" disabled aria-label="下载图表">${dlSvg}<span>下载</span></button>
  </div>
  <div class="mermaid-inner">
    <div class="mermaid-loading"><span></span><span></span><span></span></div>
  </div>
</div>`.trim()
}

function setMermaidActionsState(block: HTMLElement, payload: MermaidImagePayload | null) {
    const actionNodes = block.querySelectorAll<HTMLElement>('[data-mermaid-action]')
    actionNodes.forEach((node) => {
        const action = node.getAttribute('data-mermaid-action')
        if (action === 'zoom') {
            if (payload) {
                node.removeAttribute('disabled')
                node.removeAttribute('title')
                node.setAttribute('data-mermaid-url', payload.previewUrl)
            } else {
                node.setAttribute('disabled', 'true')
                node.removeAttribute('title')
                node.removeAttribute('data-mermaid-url')
            }
            return
        }

        if (action === 'download') {
            if (payload) {
                node.removeAttribute('disabled')
                node.setAttribute('data-mermaid-url', payload.downloadUrl || payload.previewUrl)
                node.setAttribute('data-mermaid-filename', payload.filename)
                if (payload.downloadUrl) {
                    node.removeAttribute('title')
                } else {
                    node.setAttribute('title', 'PNG will be generated on download')
                }
            } else {
                node.setAttribute('disabled', 'true')
                node.removeAttribute('data-mermaid-url')
                node.removeAttribute('title')
            }
            return
        }

        if (payload) {
            node.removeAttribute('disabled')
        } else {
            node.setAttribute('disabled', 'true')
            node.removeAttribute('data-mermaid-url')
        }
    })
}

function getSvgSize(svg: string) {
    // 优先读取 viewBox (取第三、四个值作为宽高)
    const viewBox = svg.match(/viewBox\s*=\s*["']\s*[-\d.]+\s+[-\d.]+\s+([-\d.]+)\s+([-\d.]+)\s*["']/i)
    if (viewBox) {
        const width = Number(viewBox[1])
        const height = Number(viewBox[2])
        if (Number.isFinite(width) && Number.isFinite(height) && width > 0 && height > 0) {
            return { width, height }
        }
    }

    // 次选读取数字像素属性
    const widthAttr = svg.match(/\bwidth\s*=\s*["']([\d.]+)(?:px)?["']/i)
    const heightAttr = svg.match(/\bheight\s*=\s*["']([\d.]+)(?:px)?["']/i)
    const attrW = widthAttr ? Number(widthAttr[1]) : 0
    const attrH = heightAttr ? Number(heightAttr[1]) : 0
    if (attrW > 0 && attrH > 0) {
        return { width: attrW, height: attrH }
    }

    return { width: 960, height: 540 }
}

function svgToDataUrl(svg: string) {
    return `data:image/svg+xml;charset=utf-8,${encodeURIComponent(svg)}`
}

function stripExternalSvgResources(svg: string) {
    return svg
        .replace(/<image\b[^>]*(?:href|xlink:href)\s*=\s*["']https?:\/\/[^"']*["'][^>]*>/gi, '')
        .replace(/url\(\s*["']?https?:\/\/[^"')]*["']?\s*\)/gi, 'none')
}

/** 给 SVG viewBox 四周加 padding，同步更新 width/height 并注入字体样式防止截断 */
function padSvgViewBox(svg: string, pad = 48): string {
    let result = svg
    
    // 注入显式字体样式，确保测量与渲染一致
    const styleInject = `
<style>
  svg { font-family: 'Inter', system-ui, sans-serif !important; }
  .node text, .edgeLabel text { font-family: 'Inter', system-ui, sans-serif !important; }
</style>`
    if (result.includes('</style>')) {
        result = result.replace('</style>', `  svg { font-family: 'Inter', system-ui, sans-serif !important; }\n  .node text, .edgeLabel text { font-family: 'Inter', system-ui, sans-serif !important; }\n</style>`)
    } else {
        result = result.replace(/>/, `>${styleInject}`)
    }

    // 匹配 viewBox
    const vbMatch = result.match(/viewBox\s*=\s*["']\s*([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)\s*["']/i)
    if (vbMatch) {
        const x = parseFloat(vbMatch[1]) - pad
        const y = parseFloat(vbMatch[2]) - pad
        const w = parseFloat(vbMatch[3]) + pad * 2
        const h = parseFloat(vbMatch[4]) + pad * 2
        result = result.replace(vbMatch[0], `viewBox="${x} ${y} ${w} ${h}"`)
        
        // 覆盖外层 width/height 为实际像素值
        result = result.replace(/\bwidth\s*=\s*["'][^"']*["']/i, `width="${w}"`)
        result = result.replace(/\bheight\s*=\s*["'][^"']*["']/i, `height="${h}"`)
    }
    return result
}

function getRasterDpr(width: number, height: number) {
    const rawDpr = typeof window !== 'undefined' ? window.devicePixelRatio || 1 : 1
    const baseDpr = Math.min(3, Math.max(1, rawDpr))
    const maxCanvasEdge = 4096
    const edgeLimitedDpr = maxCanvasEdge / Math.max(width, height, 1)
    return Math.max(1, Math.min(baseDpr, edgeLimitedDpr))
}

async function rasterizeSvgToPngDataUrl(svg: string, width: number, height: number): Promise<string> {
    const image = new Image()
    image.decoding = 'async'
    image.crossOrigin = 'anonymous'

    await new Promise<void>((resolve, reject) => {
        image.onload = () => resolve()
        image.onerror = () => reject(new Error('Failed to load rendered SVG'))
        image.src = svgToDataUrl(svg)
    })

    const dpr = getRasterDpr(width, height)
    const canvas = document.createElement('canvas')
    canvas.width = Math.max(1, Math.round(width * dpr))
    canvas.height = Math.max(1, Math.round(height * dpr))
    const ctx = canvas.getContext('2d')
    if (!ctx) {
        throw new Error('Canvas context unavailable')
    }

    ctx.setTransform(dpr, 0, 0, dpr, 0, 0)
    ctx.imageSmoothingEnabled = true
    ctx.imageSmoothingQuality = 'high'
    ctx.clearRect(0, 0, width, height)
    ctx.drawImage(image, 0, 0, width, height)
    return canvas.toDataURL('image/png')
}

async function svgToImageDataUrl(svg: string): Promise<MermaidImagePayload> {
    const paddedSvg = padSvgViewBox(svg)
    const fallback = getSvgSize(paddedSvg)
    const sourceWidth = fallback.width
    const sourceHeight = fallback.height

    let width = sourceWidth
    let height = sourceHeight

    const maxEdge = 2400
    const scale = Math.min(1, maxEdge / Math.max(width, height))
    width = Math.max(1, Math.round(width * scale))
    height = Math.max(1, Math.round(height * scale))

    const normalizedSvg = stripExternalSvgResources(paddedSvg)
    const candidates = normalizedSvg === paddedSvg ? [paddedSvg] : [paddedSvg, normalizedSvg]
    let lastError: unknown = null

    for (const candidate of candidates) {
        try {
            const pngUrl = await rasterizeSvgToPngDataUrl(candidate, width, height)
            return {
                previewUrl: pngUrl,
                downloadUrl: pngUrl,
                width,
                height,
                sourceWidth,
                sourceHeight,
                filename: makeMermaidFilename(),
            }
        } catch (err) {
            lastError = err
        }
    }

    const message = lastError instanceof Error ? lastError.message.toLowerCase() : String(lastError).toLowerCase()
    if (message.includes('tainted') || message.includes('security')) {
        return {
            previewUrl: svgToDataUrl(svg),
            downloadUrl: null,
            width,
            height,
            sourceWidth,
            sourceHeight,
            filename: makeMermaidFilename(),
        }
    }

    throw lastError instanceof Error ? lastError : new Error('Failed to convert Mermaid diagram to PNG')
}

function getViewportWidth() {
    const docWidth = document.documentElement?.clientWidth ?? 0
    const bodyWidth = document.body?.clientWidth ?? 0
    const winWidth = window.innerWidth ?? 0
    return Math.max(docWidth, bodyWidth, winWidth, 1)
}

function getDisplayWidthPx(payload: MermaidImagePayload) {
    const threshold = Math.max(1, Math.floor(getViewportWidth() * 0.8))
    if (payload.sourceWidth > threshold) {
        return Math.min(payload.width, threshold)
    }
    return Math.min(payload.width, payload.sourceWidth)
}

async function renderMermaidBlock(block: HTMLElement, token: number): Promise<void> {
    const tokenOnBlock = Number(block.getAttribute('data-mermaid-token') || '0')
    if (tokenOnBlock !== token) return

    const inner = block.querySelector('.mermaid-inner')
    if (!(inner instanceof HTMLElement)) return

    const encodedCode = block.getAttribute('data-mermaid-code') || ''
    const code = decodeMermaidCode(encodedCode).trim() || 'graph TD\nA-->B'

    inner.innerHTML = '<div class="mermaid-loading"><span></span><span></span><span></span></div>'
    setMermaidActionsState(block, null)

    try {
        ensureMermaid()
        const id = `mermaid-render-${++diagramCounter}`
        const { svg } = await mermaid.render(id, code)
        const imagePayload = await svgToImageDataUrl(svg)

        const latestToken = Number(block.getAttribute('data-mermaid-token') || '0')
        if (latestToken !== token) return

        const displayWidth = Math.max(1, Math.round(getDisplayWidthPx(imagePayload)))
        const displayHeight = Math.max(1, Math.round((imagePayload.height / Math.max(1, imagePayload.width)) * displayWidth))

        inner.innerHTML = `<img class="mermaid-image" src="${imagePayload.previewUrl}" alt="Mermaid diagram" width="${displayWidth}" height="${displayHeight}" style="width:${displayWidth}px;height:auto;">`
        block.setAttribute('data-mermaid-width', String(imagePayload.width))
        block.setAttribute('data-mermaid-height', String(imagePayload.height))
        block.setAttribute('data-mermaid-source-width', String(imagePayload.sourceWidth))
        block.setAttribute('data-mermaid-source-height', String(imagePayload.sourceHeight))
        block.setAttribute('data-mermaid-display-width', String(displayWidth))
        block.setAttribute('data-mermaid-url', imagePayload.previewUrl)
        block.style.removeProperty('width')
        block.style.removeProperty('max-width')
        setMermaidActionsState(block, imagePayload)
    } catch (err) {
        const latestToken = Number(block.getAttribute('data-mermaid-token') || '0')
        if (latestToken !== token) return

        const message = err instanceof Error ? err.message : String(err)
        inner.innerHTML = `<pre class="mermaid-error">Mermaid error:\n${escapeHtml(message)}</pre>`
        setMermaidActionsState(block, null)
    }
}

function scheduleMermaidRender(token: number) {
    const maxAttempts = 24
    const targetSelector = `.mermaid-block[data-mermaid-token="${token}"]`

    const run = (attempt: number) => {
        const block = document.querySelector(targetSelector)
        if (block instanceof HTMLElement) {
            void renderMermaidBlock(block, token)
            return
        }
        if (attempt >= maxAttempts) return

        if (typeof window.requestAnimationFrame === 'function') {
            window.requestAnimationFrame(() => run(attempt + 1))
        } else {
            window.setTimeout(() => run(attempt + 1), 16)
        }
    }

    run(0)
}

export function mermaidRenderPreview(
    language: string,
    content: string,
    applyPreview: (value: null | string | HTMLElement) => void,
): void | null {
    if (language !== 'mermaid') {
        return null
    }

    const code = content.trim() || 'graph TD\nA-->B'
    const token = ++renderCounter
    applyPreview(buildMermaidPreviewMarkup(code, token))
    scheduleMermaidRender(token)
}

export function refreshMermaidPreviews() {
    const blocks = document.querySelectorAll<HTMLElement>('.mermaid-block[data-mermaid-code]')
    blocks.forEach((block) => {
        const token = ++renderCounter
        block.setAttribute('data-mermaid-token', String(token))
        void renderMermaidBlock(block, token)
    })
}

export { codeBlockConfig }