TableOfContentsBlock.tsx

import React, { useEffect, useRef, useState } from "react"
import { createReactBlockSpec } from "@blocknote/react"
import { SuperedBlockNoteEditor } from "../../BlockNoteEditor"
import useReload from "../../../useReload"

export const TABLE_OF_CONTENTS_TYPE = "table-of-contents"

interface Heading {
  id: string
  level: 1 | 2 | 3 | 4
  children: Heading[]
}

type BlockText = Record<string, { id: string; text: string }>

export const TableOfContentsBlock = createReactBlockSpec(
  {
    type: TABLE_OF_CONTENTS_TYPE,
    propSchema: {},
    content: "none"
  },
  {
    parse: (el) => {
      if (el.classList.contains("mce-toc")) {
        return {}
      }

      return undefined
    },
    render: (props) => {
      const [_reload, doReload] = useReload()
      const blockText = useRef<BlockText>({})
      const editor = props.editor as unknown as SuperedBlockNoteEditor
      const [headings, setHeadings] = useState<Heading[]>([])

      const refetchAllHeadings = () => {
        let buildHeadings: Heading[] = []
        let lastHeadings: Record<number, Heading> = {}

        editor.document.forEach((block) => {
          if (block.type === "heading" && Array.isArray(block.content)) {
            const level = block.props.level
            const heading = { id: block.id, level, children: [] }
            const maybeParent = findParentHeading(lastHeadings, level)

            if (maybeParent) {
              maybeParent.children.push(heading)
            } else {
              buildHeadings.push(heading)
            }

            lastHeadings[level] = heading
            unsetHigherHeadings(lastHeadings, level)

            // This is the cleanest way I found to export a block to plaintext, but it's async and a little messy
            editor.blocksToHTMLLossy([block]).then((str) => {
              // firstChild to capture only the heading, otherwise nested content is returned
              const node = document.createRange().createContextualFragment(str).firstChild

              if (node?.textContent) {
                blockText.current[block.id] = { id: block.id, text: node.textContent }
                doReload()
              } else {
                delete blockText.current[block.id]
              }
            })
          }
        })

        setHeadings(buildHeadings)
      }

      useEffect(() => {
        refetchAllHeadings()
        return editor.onChange(refetchAllHeadings)
      }, [])

      return <RenderList editor={editor} blockText={blockText.current} headings={headings} level={1} />
    }
  }
)

function findParentHeading(headings: Record<number, Heading>, currentLevel: number) {
  // This puts skipped headings underneath the closet parent heading, which may not be desired
  for (let i = currentLevel - 1; i >= 0; i--) {
    if (headings[i]) return headings[i]
  }
}

function unsetHigherHeadings(headings: Record<number, Heading>, currentLevel: number) {
  for (let i = currentLevel + 1; i <= 10; i++) {
    delete headings[i]
  }
}

function RenderList({
  blockText,
  editor,
  headings,
  level
}: {
  blockText: BlockText
  editor: SuperedBlockNoteEditor
  headings: Heading[]
  level: number
}) {
  let listType = "list-decimal"
  if (level === 2) {
    listType = "list-[lower-alpha]"
  } else if (level === 3) {
    listType = "list-[lower-roman]"
  } else if (level === 4) {
    listType = "list-decimal"
  }

  return (
    <ul className={`${listType} pl-8`}>
      {headings.map((h) => (
        <li key={h.id}>
          <div>
            <button
              type="button"
              className="underline"
              onClick={() => {
                // TODO: Figure out how to link to anchors. Looks like others came here before and left empty handed
                // Without this, it's impossible to use an HTML export of the blocks
                // https://github.com/TypeCellOS/BlockNote/issues/478
                const element = editor.domElement.querySelector(`div[data-node-type="blockOuter"][data-id="${h.id}"]`)
                requestAnimationFrame(() => element?.scrollIntoView())
              }}
            >
              {blockText[h.id]?.text}
            </button>
          </div>
          {h.children.length > 0 && <RenderList editor={editor} blockText={blockText} headings={h.children} level={level + 1} />}
        </li>
      ))}
    </ul>
  )
}