blockflow-workbench/src/editor/extensions/SelectedBlockExtension.ts

import { Extension } from "@tiptap/core"
import { Plugin, PluginKey } from "@tiptap/pm/state"
import { Decoration, DecorationSet } from "@tiptap/pm/view"
import type { Node as ProseMirrorNode } from "@tiptap/pm/model"

export const selectedBlockPluginKey = new PluginKey<DecorationSet>("blockFlowSelectedBlock")

export const SelectedBlockExtension = Extension.create({
  name: "blockFlowSelectedBlock",

  addStorage() {
    return {
      selectedBlockId: null as string | null
    }
  },

  addProseMirrorPlugins() {
    return [
      new Plugin({
        key: selectedBlockPluginKey,
        state: {
          init: (_, state) => buildSelectedBlockDecorations(state.doc, this.storage.selectedBlockId as string | null),
          apply: (transaction, previous, _oldState, newState) => {
            const selectedBlockId = transaction.getMeta(selectedBlockPluginKey) as string | null | undefined
            if (selectedBlockId === undefined && !transaction.docChanged) {
              return previous
            }

            return buildSelectedBlockDecorations(
              newState.doc,
              selectedBlockId === undefined ? (this.storage.selectedBlockId as string | null) : selectedBlockId
            )
          }
        },
        props: {
          decorations(state) {
            return selectedBlockPluginKey.getState(state)
          }
        }
      })
    ]
  }
})

function buildSelectedBlockDecorations(doc: ProseMirrorNode, selectedBlockId: string | null): DecorationSet {
  if (selectedBlockId === null) {
    return DecorationSet.empty
  }

  const decorations: Decoration[] = []
  doc.descendants((node, position) => {
    if (node.attrs.id === selectedBlockId) {
      decorations.push(Decoration.node(position, position + node.nodeSize, { class: "bf-selected-block" }))
      return false
    }

    return true
  })

  return DecorationSet.create(doc, decorations)
}