import type { Root, ElementContent } from 'hast'
import type { Plugin } from 'unified'
import { visit } from 'unist-util-visit'

/**
 * A rehype plugin that processes text nodes and wraps text enclosed in double equal signs (==) with <mark> tags.
 * For example, "This is ==highlighted==" will be transformed into "This is <mark>highlighted</mark>".
 */
export const rehypeMark: Plugin<[], Root> = () => (tree: Root) => {
  visit(tree, 'text', (node: ElementContent, index, parent) => {
    if (!parent || typeof index !== 'number') return

    const textNode = node as ElementContent & { value: string }
    const regex = /==([^=]+)==/g
    const matches = [...textNode.value.matchAll(regex)]

    if (matches.length === 0) return

    const newNodes: ElementContent[] = []
    let lastIndex = 0

    for (const match of matches) {
      const [fullMatch, innerText] = match
      const matchStartIndex = match.index ?? 0

      if (matchStartIndex > lastIndex) {
        newNodes.push({
          type: 'text',
          value: textNode.value.slice(lastIndex, matchStartIndex),
        })
      }

      newNodes.push({
        type: 'element',
        tagName: 'mark',
        properties: {},
        children: [{ type: 'text', value: innerText }],
      })

      lastIndex = matchStartIndex + fullMatch.length
    }

    if (lastIndex < textNode.value.length) {
      newNodes.push({
        type: 'text',
        value: textNode.value.slice(lastIndex),
      })
    }

    parent.children.splice(index, 1, ...newNodes)
  })
}