/** * HAST visitor plugin that configures highlighted code blocks and adds * data-line-numbers attributes inside
 elements.
 */
import rehypePrettyCode from 'rehype-pretty-code'
import { transformerCopyButton } from '@rehype-pretty/transformers'
import { visit } from 'unist-util-visit'

const prettyCodeOptions = {
  theme: { light: 'github-light', dark: 'github-dark-dimmed' },
  keepBackground: false,
  grid: true,
  transformers: [
    transformerCopyButton({ jsx: true, visibility: 'hover', feedbackDuration: 2500 }),
  ],
}

/**
 * Rehype plugin that runs pretty-code with a React-safe copy button and then
 * adds line number indicators to code blocks.
 * @returns {import('unified').Plugin<[], import('hast').Root>}
 */
export default function codeBlockPipeline() {
  const prettyCode = rehypePrettyCode(prettyCodeOptions)
  const lineNumbers = addLineNumbers()

  return async function transformer(tree, file) {
    await prettyCode.call(this, tree, file)
    lineNumbers(tree, file)
  }
}

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

        if (!code) return

        code.properties = code.properties || {}
        code.properties['data-line-numbers'] = ''
      }
    })
  }
}