/**
 * HAST visitor plugin that adds data-line-numbers attribute to code blocks
 * inside <pre> elements.
 */
import { visit } from 'unist-util-visit'

/**
 * Rehype plugin that adds line number indicators to code blocks.
 * @returns {import('unified').Plugin<[], import('hast').Root>}
 */
export default function addLineNumbers() {
  return function attacher(tree) {
    visit(tree, 'element', function visitor(node) {
      if (
        node.tagName === 'pre' &&
        node.children?.length === 1 &&
        node.children[0].type === 'element' &&
        node.children[0].tagName === 'code'
      ) {
        const code = node.children[0]
        code.properties = code.properties || {}
        code.properties['data-line-numbers'] = ''
      }
    })
  }
}