ultisuite-client/lib/drive/extensions/docs-table-column-resizing.ts

import type { Node as ProseMirrorNode } from "@tiptap/pm/model"
import { Plugin, PluginKey, type EditorState, type Transaction } from "@tiptap/pm/state"
import {
  Decoration,
  DecorationSet,
  type EditorView,
  type NodeView,
} from "@tiptap/pm/view"
import {
  TableMap,
  TableView,
  cellAround,
  pointsAtCell,
} from "@tiptap/pm/tables"

type Dragging = { startX: number; startWidth: number }

class DocsResizeState {
  activeHandle: number
  dragging: Dragging | false

  constructor(activeHandle: number, dragging: Dragging | false) {
    this.activeHandle = activeHandle
    this.dragging = dragging
  }

  apply(tr: Transaction): DocsResizeState {
    const action = tr.getMeta(docsColumnResizingPluginKey) as
      | { setHandle?: number; setDragging?: Dragging | null }
      | undefined

    if (action && action.setHandle != null) {
      return new DocsResizeState(action.setHandle, false)
    }
    if (action && action.setDragging !== undefined) {
      return new DocsResizeState(this.activeHandle, action.setDragging ?? false)
    }
    if (this.activeHandle > -1 && tr.docChanged) {
      let handle = tr.mapping.map(this.activeHandle, -1)
      if (!pointsAtCell(tr.doc.resolve(handle))) handle = -1
      return new DocsResizeState(handle, this.dragging)
    }
    return this
  }
}

export const docsColumnResizingPluginKey = new PluginKey<DocsResizeState>(
  "docsTableColumnResizing"
)

function readDocsPageScale(view: EditorView): number {
  const stack = view.dom.closest("[data-docs-page-stack]")
  if (!stack) return 1
  const scale = Number.parseFloat(stack.getAttribute("data-docs-page-scale") ?? "1")
  return Number.isFinite(scale) && scale > 0 ? scale : 1
}

/** Keep table full width; only colgroup drives column sizes (avoids border jump on hover/drag). */
export function docsUpdateColumnsOnResize(
  node: ProseMirrorNode,
  colgroup: HTMLTableColElement,
  table: HTMLTableElement,
  defaultCellMinWidth: number,
  overrideCol?: number,
  overrideValue?: number
): void {
  let totalWidth = 0
  let fixedWidth = true
  let nextDOM = colgroup.firstChild as HTMLElement | null
  const row = node.firstChild
  if (!row) return

  for (let i = 0, col = 0; i < row.childCount; i++) {
    const { colspan, colwidth } = row.child(i).attrs as {
      colspan: number
      colwidth?: number[] | null
    }
    for (let j = 0; j < colspan; j++, col++) {
      const hasWidth = overrideCol === col ? overrideValue : colwidth && colwidth[j]
      const cssWidth = hasWidth ? `${hasWidth}px` : ""
      totalWidth += hasWidth || defaultCellMinWidth
      if (!hasWidth) fixedWidth = false
      if (!nextDOM) {
        const colEl = document.createElement("col")
        colEl.style.width = cssWidth
        colgroup.appendChild(colEl)
      } else {
        if (nextDOM.style.width !== cssWidth) nextDOM.style.width = cssWidth
        nextDOM = nextDOM.nextSibling as HTMLElement | null
      }
    }
  }

  while (nextDOM) {
    const after = nextDOM.nextSibling
    nextDOM.parentNode?.removeChild(nextDOM)
    nextDOM = after as HTMLElement | null
  }

  table.style.width = "100%"
  table.style.minWidth = fixedWidth ? "" : `${totalWidth}px`
}

export class DocsTableView extends TableView {
  constructor(node: ProseMirrorNode, defaultCellMinWidth: number, _view?: EditorView) {
    super(node, defaultCellMinWidth)
    docsUpdateColumnsOnResize(node, this.colgroup, this.table, defaultCellMinWidth)
  }

  update(node: ProseMirrorNode): boolean {
    if (node.type !== this.node.type) return false
    this.node = node
    docsUpdateColumnsOnResize(node, this.colgroup, this.table, this.defaultCellMinWidth)
    return true
  }
}

function tableNodeName(state: EditorState): string {
  for (const name of Object.keys(state.schema.nodes)) {
    if (state.schema.nodes[name].spec.tableRole === "table") return name
  }
  return "table"
}

function domCellAround(target: EventTarget | null): HTMLElement | null {
  let el = target as HTMLElement | null
  while (el && el.nodeName !== "TD" && el.nodeName !== "TH") {
    el =
      el.classList?.contains("ProseMirror") ? null : (el.parentNode as HTMLElement | null)
  }
  return el
}

function edgeCell(
  view: EditorView,
  event: MouseEvent,
  side: "left" | "right",
  handleWidth: number
): number {
  const offset = side === "right" ? -handleWidth : handleWidth
  const found = view.posAtCoords({
    left: event.clientX + offset,
    top: event.clientY,
  })
  if (!found) return -1
  const { pos } = found
  const $cell = cellAround(view.state.doc.resolve(pos))
  if (!$cell) return -1
  if (side === "right") return $cell.pos
  const map = TableMap.get($cell.node(-1))
  const start = $cell.start(-1)
  const index = map.map.indexOf($cell.pos - start)
  return index % map.width === 0 ? -1 : start + map.map[index - 1]
}

function currentColWidth(
  view: EditorView,
  cellPos: number,
  { colspan, colwidth }: { colspan: number; colwidth?: number[] | null }
): number {
  const stored = colwidth && colwidth[colwidth.length - 1]
  if (stored) return stored
  const dom = view.domAtPos(cellPos)
  const cellDom = dom.node.childNodes[dom.offset] as HTMLElement
  let domWidth = cellDom.offsetWidth
  let parts = colspan
  if (colwidth) {
    for (let i = 0; i < colspan; i++) {
      if (colwidth[i]) {
        domWidth -= colwidth[i]!
        parts--
      }
    }
  }
  return domWidth / parts
}

function draggedWidth(dragging: Dragging, event: MouseEvent, minWidth: number, scale: number) {
  const offset = (event.clientX - dragging.startX) / scale
  return Math.max(minWidth, dragging.startWidth + offset)
}

function updateHandle(view: EditorView, value: number) {
  view.dispatch(view.state.tr.setMeta(docsColumnResizingPluginKey, { setHandle: value }))
}

function displayColumnWidth(
  view: EditorView,
  cell: number,
  width: number,
  defaultCellMinWidth: number
) {
  const $cell = view.state.doc.resolve(cell)
  const table = $cell.node(-1)
  const start = $cell.start(-1)
  const col =
    TableMap.get(table).colCount($cell.pos - start) + ($cell.nodeAfter?.attrs.colspan ?? 1) - 1
  let dom: Node | null = view.domAtPos($cell.start(-1)).node
  while (dom && dom.nodeName !== "TABLE") dom = dom.parentNode
  if (!dom || dom.nodeName !== "TABLE") return
  docsUpdateColumnsOnResize(
    table,
    (dom as HTMLTableElement).firstChild as HTMLTableColElement,
    dom as HTMLTableElement,
    defaultCellMinWidth,
    col,
    width
  )
}

function zeroes(n: number) {
  return Array<number>(n).fill(0)
}

function updateColumnWidth(view: EditorView, cell: number, width: number) {
  const $cell = view.state.doc.resolve(cell)
  const table = $cell.node(-1)
  const map = TableMap.get(table)
  const start = $cell.start(-1)
  const col = map.colCount($cell.pos - start) + ($cell.nodeAfter?.attrs.colspan ?? 1) - 1
  const tr = view.state.tr
  for (let row = 0; row < map.height; row++) {
    const mapIndex = row * map.width + col
    if (row && map.map[mapIndex] === map.map[mapIndex - map.width]) continue
    const pos = map.map[mapIndex]
    const attrs = table.nodeAt(pos)!.attrs
    const index = attrs.colspan === 1 ? 0 : col - map.colCount(pos)
    if (attrs.colwidth && attrs.colwidth[index] === width) continue
    const colwidth = attrs.colwidth ? attrs.colwidth.slice() : zeroes(attrs.colspan)
    colwidth[index] = width
    tr.setNodeMarkup(start + pos, null, { ...attrs, colwidth })
  }
  if (tr.docChanged) view.dispatch(tr)
}

function handleDecorations(state: EditorState, cell: number) {
  const decorations: Decoration[] = []
  const $cell = state.doc.resolve(cell)
  const table = $cell.node(-1)
  if (!table) return DecorationSet.empty
  const map = TableMap.get(table)
  const start = $cell.start(-1)
  const col =
    map.colCount($cell.pos - start) + ($cell.nodeAfter?.attrs.colspan ?? 1) - 1
  const pluginState = docsColumnResizingPluginKey.getState(state)

  for (let row = 0; row < map.height; row++) {
    const index = col + row * map.width
    if (
      (col === map.width - 1 || map.map[index] !== map.map[index + 1]) &&
      (row === 0 || map.map[index] !== map.map[index - map.width])
    ) {
      const cellPos = map.map[index]
      const pos = start + cellPos + table.nodeAt(cellPos)!.nodeSize - 1
      const dom = document.createElement("div")
      dom.className = "column-resize-handle"
      if (pluginState?.dragging) {
        decorations.push(
          Decoration.node(start + cellPos, start + cellPos + table.nodeAt(cellPos)!.nodeSize, {
            class: "column-resize-dragging",
          })
        )
      }
      decorations.push(Decoration.widget(pos, dom))
    }
  }
  return DecorationSet.create(state.doc, decorations)
}

function handleMouseMove(
  view: EditorView,
  event: MouseEvent,
  handleWidth: number,
  lastColumnResizable: boolean
) {
  if (!view.editable) return
  const pluginState = docsColumnResizingPluginKey.getState(view.state)
  if (!pluginState || pluginState.dragging) return

  const target = domCellAround(event.target)
  let cell = -1
  if (target) {
    const { left, right } = target.getBoundingClientRect()
    if (event.clientX - left <= handleWidth) cell = edgeCell(view, event, "left", handleWidth)
    else if (right - event.clientX <= handleWidth) cell = edgeCell(view, event, "right", handleWidth)
  }

  if (cell === pluginState.activeHandle) return

  if (!lastColumnResizable && cell !== -1) {
    const $cell = view.state.doc.resolve(cell)
    const table = $cell.node(-1)
    const map = TableMap.get(table)
    const tableStart = $cell.start(-1)
    if (
      map.colCount($cell.pos - tableStart) + ($cell.nodeAfter?.attrs.colspan ?? 1) - 1 ===
      map.width - 1
    ) {
      return
    }
  }
  updateHandle(view, cell)
}

function handleMouseLeave(view: EditorView) {
  if (!view.editable) return
  const pluginState = docsColumnResizingPluginKey.getState(view.state)
  if (pluginState && pluginState.activeHandle > -1 && !pluginState.dragging) {
    updateHandle(view, -1)
  }
}

function handleMouseDown(
  view: EditorView,
  event: MouseEvent,
  cellMinWidth: number,
  defaultCellMinWidth: number
) {
  if (!view.editable) return false
  const win = view.dom.ownerDocument.defaultView ?? window
  const pluginState = docsColumnResizingPluginKey.getState(view.state)
  if (!pluginState || pluginState.activeHandle === -1 || pluginState.dragging) return false

  const cell = view.state.doc.nodeAt(pluginState.activeHandle)
  if (!cell) return false
  const width = currentColWidth(view, pluginState.activeHandle, {
    colspan: cell.attrs.colspan as number,
    colwidth: cell.attrs.colwidth as number[] | null | undefined,
  })
  const scale = readDocsPageScale(view)

  view.dispatch(
    view.state.tr.setMeta(docsColumnResizingPluginKey, {
      setDragging: { startX: event.clientX, startWidth: width },
    })
  )

  const finish = (ev: MouseEvent) => {
    win.removeEventListener("mouseup", finish)
    win.removeEventListener("mousemove", move)
    const state = docsColumnResizingPluginKey.getState(view.state)
    if (state?.dragging) {
      updateColumnWidth(
        view,
        state.activeHandle,
        draggedWidth(state.dragging, ev, cellMinWidth, scale)
      )
      view.dispatch(view.state.tr.setMeta(docsColumnResizingPluginKey, { setDragging: null }))
    }
  }

  const move = (ev: MouseEvent) => {
    if (!ev.buttons) return finish(ev)
    const state = docsColumnResizingPluginKey.getState(view.state)
    if (!state?.dragging) return
    displayColumnWidth(
      view,
      state.activeHandle,
      draggedWidth(state.dragging, ev, cellMinWidth, scale),
      defaultCellMinWidth
    )
  }

  displayColumnWidth(view, pluginState.activeHandle, width, defaultCellMinWidth)
  win.addEventListener("mouseup", finish)
  win.addEventListener("mousemove", move)
  event.preventDefault()
  return true
}

export function docsColumnResizing({
  handleWidth = 5,
  cellMinWidth = 25,
  defaultCellMinWidth = 100,
  View = DocsTableView,
  lastColumnResizable = true,
}: {
  handleWidth?: number
  cellMinWidth?: number
  defaultCellMinWidth?: number
  View?: (new (node: ProseMirrorNode, cellMinWidth: number, view: EditorView) => NodeView) | null
  lastColumnResizable?: boolean
} = {}) {
  const plugin = new Plugin<DocsResizeState>({
    key: docsColumnResizingPluginKey,
    state: {
      init(_, state) {
        const tableName = tableNodeName(state)
        if (View && plugin.spec.props?.nodeViews) {
          plugin.spec.props.nodeViews[tableName] = (node, view) =>
            new View(node, defaultCellMinWidth, view)
        }
        return new DocsResizeState(-1, false)
      },
      apply(tr, prev) {
        return prev.apply(tr)
      },
    },
    props: {
      attributes(state) {
        const pluginState = docsColumnResizingPluginKey.getState(state)
        return pluginState && pluginState.activeHandle > -1 ? { class: "resize-cursor" } : {}
      },
      handleDOMEvents: {
        mousemove: (view, event) => {
          handleMouseMove(view, event, handleWidth, lastColumnResizable)
        },
        mouseleave: (view) => {
          handleMouseLeave(view)
        },
        mousedown: (view, event) => {
          return handleMouseDown(view, event, cellMinWidth, defaultCellMinWidth)
        },
      },
      decorations(state) {
        const pluginState = docsColumnResizingPluginKey.getState(state)
        if (pluginState && pluginState.activeHandle > -1) {
          return handleDecorations(state, pluginState.activeHandle)
        }
        return undefined
      },
      nodeViews: {},
    },
  })
  return plugin
}