1import { mdxAnnotations } from 'mdx-annotations'
  2import { visit } from 'unist-util-visit'
  3import rehypeMdxTitle from 'rehype-mdx-title'
  4import shiki from 'shiki'
  5import { toString } from 'mdast-util-to-string'
  6import * as acorn from 'acorn'
  7import { slugifyWithCounter } from '@sindresorhus/slugify'
  8
  9function rehypeParseCodeBlocks() {
 10  return (tree) => {
 11    visit(tree, 'element', (node, _nodeIndex, parentNode) => {
 12      if (node.tagName === 'code' && node.properties.className) {
 13        parentNode.properties.language = node.properties.className[0]?.replace(
 14          /^language-/,
 15          ''
 16        )
 17      }
 18    })
 19  }
 20}
 21
 22let highlighter
 23
 24function rehypeShiki() {
 25  return async (tree) => {
 26    highlighter =
 27      highlighter ?? (await shiki.getHighlighter({ theme: 'css-variables' }))
 28
 29    visit(tree, 'element', (node) => {
 30      if (node.tagName === 'pre' && node.children[0]?.tagName === 'code') {
 31        let codeNode = node.children[0]
 32        let textNode = codeNode.children[0]
 33
 34        node.properties.code = textNode.value
 35
 36        if (node.properties.language) {
 37          let tokens = highlighter.codeToThemedTokens(
 38            textNode.value,
 39            node.properties.language
 40          )
 41
 42          textNode.value = shiki.renderToHtml(tokens, {
 43            elements: {
 44              pre: ({ children }) => children,
 45              code: ({ children }) => children,
 46              line: ({ children }) => `<span>${children}</span>`,
 47            },
 48          })
 49        }
 50      }
 51    })
 52  }
 53}
 54
 55function rehypeSlugify() {
 56  return (tree) => {
 57    let slugify = slugifyWithCounter()
 58    visit(tree, 'element', (node) => {
 59      if (node.tagName === 'h2' && !node.properties.id) {
 60        node.properties.id = slugify(toString(node))
 61      }
 62    })
 63  }
 64}
 65
 66function rehypeAddMDXExports(getExports) {
 67  return (tree) => {
 68    let exports = Object.entries(getExports(tree))
 69
 70    for (let [name, value] of exports) {
 71      for (let node of tree.children) {
 72        if (
 73          node.type === 'mdxjsEsm' &&
 74          new RegExp(`export\\s+const\\s+${name}\\s*=`).test(node.value)
 75        ) {
 76          return
 77        }
 78      }
 79
 80      let exportStr = `export const ${name} = ${value}`
 81
 82      tree.children.push({
 83        type: 'mdxjsEsm',
 84        value: exportStr,
 85        data: {
 86          estree: acorn.parse(exportStr, {
 87            sourceType: 'module',
 88            ecmaVersion: 'latest',
 89          }),
 90        },
 91      })
 92    }
 93  }
 94}
 95
 96function getSections(node) {
 97  let sections = []
 98
 99  for (let child of node.children ?? []) {
100    if (child.type === 'element' && child.tagName === 'h2') {
101      sections.push(`{
102        title: ${JSON.stringify(toString(child))},
103        id: ${JSON.stringify(child.properties.id)},
104        ...${child.properties.annotation}
105      }`)
106    } else if (child.children) {
107      sections.push(...getSections(child))
108    }
109  }
110
111  return sections
112}
113
114export const rehypePlugins = [
115  mdxAnnotations.rehype,
116  rehypeParseCodeBlocks,
117  rehypeShiki,
118  rehypeSlugify,
119  rehypeMdxTitle,
120  [
121    rehypeAddMDXExports,
122    (tree) => ({
123      sections: `[${getSections(tree).join()}]`,
124    }),
125  ],
126]