import type { Root, Element } from 'hast'
import type { Plugin } from 'unified'
import { visit } from 'unist-util-visit'
import { toText } from 'hast-util-to-text'

/**
 * A rehype plugin that processes code blocks and adds metadata.
 * Specifically, it looks for <pre> elements containing <code> elements,
 * extracts the language from the className, and adds a data-language attribute to the <pre> element.
 * If the language is "mermaid", it transforms the <pre> into a <div> with a data-mermaid attribute containing the code content.
 */
export const rehypeCodeMeta: Plugin<[], Root> = () => (tree: Root) => {
  visit(tree, 'element', (node: Element) => {
    if (node.tagName !== 'pre') return
    const codeNode = node.children.find(
      (c): c is Element => c.type === 'element' && c.tagName === 'code',
    )
    if (!codeNode) return

    const classNames = (codeNode.properties?.className ?? []) as string[]
    const language =
      classNames.find((c) => c.startsWith('language-'))?.replace('language-', '') ?? 'text'

    if (language === 'mermaid') {
      node.tagName = 'div'
      node.properties = { 'data-mermaid': toText(codeNode) }
      node.children = []
      return
    }

    node.properties = { ...node.properties, 'data-language': language }
  })
}