import type { Root, Element } from 'hast'
import type { Plugin } from 'unified'
import { visit } from 'unist-util-visit'
import { fromHtmlIsomorphic } from 'hast-util-from-html-isomorphic'
import { renderMermaidToSvg } from './mermaid-render'

/**
 * A rehype plugin that processes <div> elements with a data-mermaid attribute and renders the Mermaid chart to SVG.
 * It replaces the <div> with the rendered SVG or displays an error message if rendering fails.
 */
export const rehypeMermaidSsg: Plugin<[], Root> = () => async (tree: Root) => {
  const targets: { node: Element; parent: Element | Root; index: number }[] = []

  visit(tree, 'element', (node: Element, index, parent) => {
    if (
      node.tagName === 'div' &&
      node.properties?.['data-mermaid'] &&
      parent &&
      typeof index === 'number'
    ) {
      targets.push({ node, parent: parent as Element | Root, index })
    }
  })

  await Promise.all(
    targets.map(async ({ node, parent, index }) => {
      const chart = String(node.properties!['data-mermaid'])
      try {
        const svg = await renderMermaidToSvg(chart)
        const svgTree = fromHtmlIsomorphic(svg, { fragment: true })
        const svgNode = svgTree.children.find((c) => c.type === 'element') as Element | undefined
        if (svgNode) {
          ;(parent.children as Element[])[index] = svgNode
        }
      } catch (err) {
        node.tagName = 'pre'
        node.properties = { className: ['text-red-500', 'text-sm'] }
        node.children = [{ type: 'text', value: `Mermaid render error: ${(err as Error).message}` }]
      }
    }),
  )
}