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

import { mergeAttributes } from "@tiptap/core"
import { Table } from "@tiptap/extension-table"
import { TableCell } from "@tiptap/extension-table-cell"
import { TableHeader } from "@tiptap/extension-table-header"
import { TableRow } from "@tiptap/extension-table-row"
import { tableEditing } from "@tiptap/pm/tables"
import {
  buildDocsTableCellStyle,
  buildDocsTableRowStyle,
  buildDocsTableStyle,
} from "@/lib/drive/docs-table-styles"
import {
  DocsTableView,
  docsColumnResizing,
} from "@/lib/drive/extensions/docs-table-column-resizing"

function parseJsonAttr(raw: string | null): unknown {
  if (!raw) return null
  try {
    return JSON.parse(raw)
  } catch {
    return null
  }
}

function borderAttribute(name: string) {
  return {
    default: null,
    parseHTML: (element: HTMLElement) => {
      const fromData = parseJsonAttr(element.getAttribute(`data-${name}`))
      if (fromData) return fromData
      const cssName = name.replace(/([A-Z])/g, "-$1").toLowerCase()
      const fromStyle = element.style.getPropertyValue(cssName)
      return fromStyle || null
    },
    renderHTML: (attributes: Record<string, unknown>) => {
      const value = attributes[name]
      if (!value) return {}
      return { [`data-${name}`]: JSON.stringify(value) }
    },
  }
}

const sharedCellAttributes = () => ({
  backgroundColor: {
    default: null,
    parseHTML: (element: HTMLElement) =>
      element.getAttribute("data-background-color") ||
      element.style.backgroundColor ||
      null,
    renderHTML: (attributes: Record<string, unknown>) => {
      if (!attributes.backgroundColor) return {}
      return { "data-background-color": String(attributes.backgroundColor) }
    },
  },
  verticalAlign: {
    default: null,
    parseHTML: (element: HTMLElement) =>
      element.getAttribute("data-vertical-align") ||
      element.style.verticalAlign ||
      null,
    renderHTML: (attributes: Record<string, unknown>) => {
      if (!attributes.verticalAlign) return {}
      return { "data-vertical-align": String(attributes.verticalAlign) }
    },
  },
  borderTop: borderAttribute("borderTop"),
  borderRight: borderAttribute("borderRight"),
  borderBottom: borderAttribute("borderBottom"),
  borderLeft: borderAttribute("borderLeft"),
})

export const DocsTableCell = TableCell.extend({
  renderHTML({ node, HTMLAttributes }) {
    const style = buildDocsTableCellStyle(node.attrs)
    return [
      "td",
      mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, style ? { style } : {}),
      0,
    ]
  },
  addAttributes() {
    return {
      ...this.parent?.(),
      ...sharedCellAttributes(),
    }
  },
})

export const DocsTableHeader = TableHeader.extend({
  renderHTML({ node, HTMLAttributes }) {
    const style = buildDocsTableCellStyle(node.attrs)
    return [
      "th",
      mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, style ? { style } : {}),
      0,
    ]
  },
  addAttributes() {
    return {
      ...this.parent?.(),
      ...sharedCellAttributes(),
    }
  },
})

export const DocsTableRow = TableRow.extend({
  addAttributes() {
    return {
      ...this.parent?.(),
      rowHeight: {
        default: null,
        parseHTML: (element: HTMLElement) =>
          element.getAttribute("data-row-height") || element.style.height || null,
        renderHTML: (attributes: Record<string, unknown>) => {
          if (!attributes.rowHeight) return {}
          return { "data-row-height": String(attributes.rowHeight) }
        },
      },
      rowHeightRule: {
        default: null,
        parseHTML: (element: HTMLElement) => element.getAttribute("data-row-height-rule"),
        renderHTML: (attributes: Record<string, unknown>) => {
          if (!attributes.rowHeightRule) return {}
          return { "data-row-height-rule": String(attributes.rowHeightRule) }
        },
      },
      header: {
        default: false,
        parseHTML: (element: HTMLElement) => element.getAttribute("data-header-row") === "true",
        renderHTML: (attributes: Record<string, unknown>) => {
          if (!attributes.header) return {}
          return { "data-header-row": "true" }
        },
      },
    }
  },
  renderHTML({ node, HTMLAttributes }) {
    const style = buildDocsTableRowStyle(node.attrs)
    return [
      "tr",
      mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, style ? { style } : {}),
      0,
    ]
  },
})

export const DocsTable = Table.extend({
  addProseMirrorPlugins() {
    const isResizable = this.options.resizable && this.editor.isEditable
    return [
      ...(isResizable
        ? [
            docsColumnResizing({
              handleWidth: this.options.handleWidth,
              cellMinWidth: this.options.cellMinWidth,
              defaultCellMinWidth: this.options.cellMinWidth,
              View: DocsTableView,
              lastColumnResizable: this.options.lastColumnResizable,
            }),
          ]
        : []),
      tableEditing({
        allowTableNodeSelection: this.options.allowTableNodeSelection,
      }),
    ]
  },
  addAttributes() {
    return {
      alignment: {
        default: null,
        parseHTML: (element: HTMLElement) => element.getAttribute("data-alignment"),
        renderHTML: (attributes: Record<string, unknown>) => {
          if (!attributes.alignment) return {}
          return { "data-alignment": String(attributes.alignment) }
        },
      },
      layout: {
        default: null,
        parseHTML: (element: HTMLElement) => element.getAttribute("data-layout"),
        renderHTML: (attributes: Record<string, unknown>) => {
          if (!attributes.layout) return {}
          return { "data-layout": String(attributes.layout) }
        },
      },
      cellSpacing: {
        default: null,
        parseHTML: (element: HTMLElement) => {
          const raw = element.getAttribute("data-cell-spacing")
          return raw ? Number(raw) : null
        },
        renderHTML: (attributes: Record<string, unknown>) => {
          if (attributes.cellSpacing == null) return {}
          return { "data-cell-spacing": String(attributes.cellSpacing) }
        },
      },
    }
  },
  renderHTML({ node, HTMLAttributes }) {
    const style = buildDocsTableStyle(node.attrs)
    return [
      "table",
      mergeAttributes(this.options.HTMLAttributes, HTMLAttributes, style ? { style } : {}),
      0,
    ]
  },
})

export const DOCS_TABLE_EXTENSION_OPTIONS = {
  resizable: true,
  renderWrapper: true,
  cellMinWidth: 48,
  handleWidth: 6,
  lastColumnResizable: true,
  allowTableNodeSelection: true,
  HTMLAttributes: {
    class: "docs-table",
  },
} as const