Momento/memento-note/components/tiptap-columns-extension.tsx

'use client'

import { Node, mergeAttributes, findParentNode } from '@tiptap/core'
import { TextSelection } from '@tiptap/pm/state'
import type { EditorState, Transaction } from '@tiptap/pm/state'
import { Node as PMNode } from '@tiptap/pm/model'

export const ColumnNode = Node.create({
  name: 'column',
  content: 'block+',
  isolating: true,
  defining: true,

  addAttributes() {
    return {
      index: {
        default: 0,
        parseHTML: (el) => Number(el.getAttribute('index')) || 0,
        renderHTML: (attrs) => ({ index: attrs.index }),
      },
    }
  },

  parseHTML() {
    return [{ tag: 'div[data-type="column"]' }]
  },

  renderHTML({ HTMLAttributes }) {
    return [
      'div',
      mergeAttributes(HTMLAttributes, {
        'data-type': 'column',
        class: 'tiptap-column',
      }),
      0,
    ]
  },
})

export const ColumnsExtension = Node.create({
  name: 'columns',
  group: 'block',
  content: 'column{1,}',
  isolating: true,
  defining: true,
  allowGapCursor: false,

  addAttributes() {
    return {
      cols: {
        default: 2,
        parseHTML: (el) => Number(el.getAttribute('cols')) || 2,
        renderHTML: (attrs) => ({ cols: attrs.cols }),
      },
    }
  },

  parseHTML() {
    return [{ tag: 'div[data-type="columns"]' }]
  },

  renderHTML({ HTMLAttributes }) {
    return [
      'div',
      mergeAttributes(HTMLAttributes, {
        'data-type': 'columns',
        class: 'tiptap-columns',
      }),
      0,
    ]
  },

  addCommands() {
    return {
      addColumnAfter: () => ({ state, dispatch }: { state: EditorState; dispatch?: (tr: Transaction) => void }) => {
        return addOrDeleteCol(state, dispatch, 'addAfter')
      },
      addColumnBefore: () => ({ state, dispatch }: { state: EditorState; dispatch?: (tr: Transaction) => void }) => {
        return addOrDeleteCol(state, dispatch, 'addBefore')
      },
      deleteColumn: () => ({ state, dispatch }: { state: EditorState; dispatch?: (tr: Transaction) => void }) => {
        return addOrDeleteCol(state, dispatch, 'delete')
      },
    }
  },

  addKeyboardShortcuts() {
    return {
      'Mod-Shift-L': () => this.editor.commands.insertContent({
        type: this.name,
        attrs: { cols: 2 },
        content: [
          { type: 'column', attrs: { index: 0 }, content: [{ type: 'paragraph' }] },
          { type: 'column', attrs: { index: 1 }, content: [{ type: 'paragraph' }] },
        ],
      }),
    }
  },
})

function addOrDeleteCol(state: EditorState, dispatch?: (tr: Transaction) => void, type?: 'addBefore' | 'addAfter' | 'delete'): boolean {
  const maybeColumns = findParentNode((node: PMNode) => node.type.name === 'columns')(state.selection)
  const maybeColumn = findParentNode((node: PMNode) => node.type.name === 'column')(state.selection)

  if (!maybeColumns || !maybeColumn) return false

  const cols = maybeColumns.node
  const colIndex = maybeColumn.node.attrs.index
  const colsJSON = cols.toJSON()

  let nextIndex = colIndex

  if (type === 'delete') {
    if (colsJSON.content.length <= 1) return false
    nextIndex = Math.max(0, colIndex - 1)
    colsJSON.content.splice(colIndex, 1)
  } else {
    nextIndex = type === 'addBefore' ? colIndex : colIndex + 1
    colsJSON.content.splice(nextIndex, 0, {
      type: 'column',
      attrs: { index: colIndex },
      content: [{ type: 'paragraph' }],
    })
  }

  colsJSON.attrs.cols = colsJSON.content.length
  colsJSON.content.forEach((colJSON: any, index: number) => {
    colJSON.attrs.index = index
  })

  const nextCols = PMNode.fromJSON(state.schema, colsJSON)

  let nextSelectPos = maybeColumns.pos + 1
  nextCols.content.forEach((col: PMNode, _pos: number, index: number) => {
    if (index < nextIndex) nextSelectPos += col.nodeSize
  })

  if (dispatch) {
    const tr = state.tr
    tr.replaceWith(maybeColumns.pos, maybeColumns.pos + maybeColumns.node.nodeSize, nextCols)
    tr.setSelection(TextSelection.near(tr.doc.resolve(nextSelectPos + 1)))
    dispatch(tr)
  }

  return true
}

export function insertColumnsBlock(editor: any, columns: number = 2) {
  const content = Array.from({ length: columns }, (_, i) => ({
    type: 'column',
    attrs: { index: i },
    content: [{ type: 'paragraph' }],
  }))
  editor.chain().focus().insertContent({
    type: 'columns',
    attrs: { cols: columns },
    content,
  }).run()
}