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

function isBlank(node: ElementContent) {
  return (
    (node.type === 'text' && node.value.trim() === '') ||
    (node.type === 'element' && node.tagName === 'br')
  )
}

function isImage(node: ElementContent): node is Element {
  return node.type === 'element' && node.tagName === 'img'
}

/**
 * A rehype plugin that lifts <img> elements out of <p> tags, since images render as <figure>,
 * which can't be nested inside a paragraph.
 *
 * A paragraph containing only an image is replaced by the image. A paragraph mixing text and
 * images (e.g. "**Screenshots:**" followed by images on the next lines) is split: each image
 * becomes a standalone node and the surrounding inline content stays in its own <p>.
 */
export const rehypeUnwrapImages: Plugin<[], Root> = () => (tree: Root) => {
  visit(tree, 'element', (node: Element, index, parent) => {
    if (node.tagName !== 'p' || !parent || typeof index !== 'number') return
    if (!node.children.some(isImage)) return

    const replacement: ElementContent[] = []
    let run: ElementContent[] = []

    const flushRun = () => {
      // Drop whitespace and line breaks left over around the lifted images
      if (run.some((child) => !isBlank(child))) {
        replacement.push({ ...node, children: trimBlankEdges(run) })
      }
      run = []
    }

    for (const child of node.children) {
      if (isImage(child)) {
        flushRun()
        replacement.push(child)
      } else {
        run.push(child)
      }
    }
    flushRun()

    parent.children.splice(index, 1, ...replacement)
    return [SKIP, index + replacement.length]
  })
}

function trimBlankEdges(nodes: ElementContent[]): ElementContent[] {
  let start = 0
  let end = nodes.length
  while (start < end && isBlank(nodes[start])) start++
  while (end > start && isBlank(nodes[end - 1])) end--
  return nodes.slice(start, end)
}