import { TableNode } from '@lexical/table'

import {
  DOMConversionMap,
  DOMExportOutput,
  EditorConfig,
  LexicalEditor,
  NodeKey,
  SerializedElementNode
} from 'lexical'

export type SerializedTableNode = SerializedElementNode & {
  type: 'custom-table'
  version: 1
}

export class CustomTableNode extends TableNode {
  static getType(): string {
    return 'custom-table'
  }
  constructor(key?: NodeKey) {
    super(key)
  }
  static clone(node: TableNode): CustomTableNode {
    return new CustomTableNode(node.__key)
  }
  createDOM(config: EditorConfig, editor?: LexicalEditor): HTMLElement {
    return super.createDOM(config, editor)
  }
  static importDOM(): DOMConversionMap | null {
    return super.importDOM()
  }
  static importJSON(_serializedNode: SerializedTableNode): CustomTableNode {
    return new CustomTableNode()
  }
  exportJSON(): SerializedTableNode {
    return {
      ...super.exportJSON(),
      type: 'custom-table',
      version: 1
    }
  }
  exportDOM(editor: LexicalEditor): DOMExportOutput {
    const { element, after } = super.exportDOM(editor)
    return {
      element,
      after(generatedElement) {
        const tableElement = after?.(generatedElement) || ''
        const container = document.createElement('div')
        container.classList.add('editor__table-container')
        container.append(tableElement)
        return container
      }
    }
  }
}