From 48206cbc36181adb10b6be72275d985a224359d1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89mile=20R=C3=A9?= <emile@getprobo.com>
Date: Wed, 25 Mar 2026 15:32:29 +0400
Subject: [PATCH] Add header row / column option
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Émile Ré <emile@getprobo.com>
---
 .../ui/src/RichEditor/TableColumnMenu.tsx     | 55 +++++++++++++++++++
 packages/ui/src/RichEditor/TableRowMenu.tsx   | 55 +++++++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/packages/ui/src/RichEditor/TableColumnMenu.tsx b/packages/ui/src/RichEditor/TableColumnMenu.tsx
index 8d7f398a4..831cce62f 100644
--- a/packages/ui/src/RichEditor/TableColumnMenu.tsx
+++ b/packages/ui/src/RichEditor/TableColumnMenu.tsx
@@ -16,6 +16,7 @@ import {
 import {
   BroomIcon,
   CopyIcon,
+  CrownSimpleIcon,
   DotsThreeIcon,
   PlusIcon,
   TrashIcon,
@@ -453,6 +454,54 @@ export function TableColumnMenu({ editor }: TableColumnMenuProps) {
 
   const currentCol = hoveredCol;
 
+  const isFirstColumn = currentCol?.colIndex === 0;
+
+  const isHeaderColumn = (): boolean => {
+    if (!currentCol || currentCol.colIndex !== 0) return false;
+    try {
+      const table = editor.state.doc.nodeAt(currentCol.tableStart - 1);
+      if (!table) return false;
+      const map = TableMap.get(table);
+      for (let row = 0; row < map.height; row++) {
+        const cellPos = map.map[row * map.width] + currentCol.tableStart;
+        const cellNode = editor.state.doc.nodeAt(cellPos);
+        if (!cellNode || cellNode.type.name !== "tableHeader") return false;
+      }
+      return true;
+    } catch {
+      return false;
+    }
+  };
+
+  const handleToggleHeaderColumn = () => {
+    if (!currentCol || currentCol.colIndex !== 0) return;
+    const { tableStart } = currentCol;
+
+    try {
+      const table = editor.state.doc.nodeAt(tableStart - 1);
+      if (!table) return;
+
+      const map = TableMap.get(table);
+      const { tr, schema } = editor.state;
+      const targetType = isHeaderColumn()
+        ? schema.nodes.tableCell
+        : schema.nodes.tableHeader;
+
+      for (let row = 0; row < map.height; row++) {
+        const cellPos = map.map[row * map.width] + tableStart;
+        const cellNode = editor.state.doc.nodeAt(cellPos);
+        if (!cellNode) continue;
+        tr.setNodeMarkup(cellPos, targetType, cellNode.attrs);
+      }
+
+      editor.view.dispatch(tr);
+    } catch {
+      // table may have changed
+    }
+
+    setMenuOpen(false);
+  };
+
   const handleDeleteColumn = () => {
     if (!currentCol) return;
     const { colIndex, tableStart } = currentCol;
@@ -634,6 +683,12 @@ export function TableColumnMenu({ editor }: TableColumnMenuProps) {
           onMouseDown={e => e.preventDefault()}
           className={menu()}
         >
+          {isFirstColumn && (
+            <MenuButton active={isHeaderColumn()} onClick={handleToggleHeaderColumn}>
+              <CrownSimpleIcon size={16} weight="bold" />
+              Header column
+            </MenuButton>
+          )}
           <MenuButton onClick={handleInsertLeft}>
             <PlusIcon size={16} weight="bold" />
             Insert column left
diff --git a/packages/ui/src/RichEditor/TableRowMenu.tsx b/packages/ui/src/RichEditor/TableRowMenu.tsx
index 9ddd1b433..cbaa13436 100644
--- a/packages/ui/src/RichEditor/TableRowMenu.tsx
+++ b/packages/ui/src/RichEditor/TableRowMenu.tsx
@@ -16,6 +16,7 @@ import {
 import {
   BroomIcon,
   CopyIcon,
+  CrownSimpleIcon,
   DotsThreeVerticalIcon,
   PlusIcon,
   TrashIcon,
@@ -449,6 +450,54 @@ export function TableRowMenu({ editor }: TableRowMenuProps) {
 
   const currentRow = hoveredRow;
 
+  const isFirstRow = currentRow?.rowIndex === 0;
+
+  const isHeaderRow = (): boolean => {
+    if (!currentRow || currentRow.rowIndex !== 0) return false;
+    try {
+      const table = editor.state.doc.nodeAt(currentRow.tableStart - 1);
+      if (!table) return false;
+      const map = TableMap.get(table);
+      for (let col = 0; col < map.width; col++) {
+        const cellPos = map.map[col] + currentRow.tableStart;
+        const cellNode = editor.state.doc.nodeAt(cellPos);
+        if (!cellNode || cellNode.type.name !== "tableHeader") return false;
+      }
+      return true;
+    } catch {
+      return false;
+    }
+  };
+
+  const handleToggleHeaderRow = () => {
+    if (!currentRow || currentRow.rowIndex !== 0) return;
+    const { tableStart } = currentRow;
+
+    try {
+      const table = editor.state.doc.nodeAt(tableStart - 1);
+      if (!table) return;
+
+      const map = TableMap.get(table);
+      const { tr, schema } = editor.state;
+      const targetType = isHeaderRow()
+        ? schema.nodes.tableCell
+        : schema.nodes.tableHeader;
+
+      for (let col = 0; col < map.width; col++) {
+        const cellPos = map.map[col] + tableStart;
+        const cellNode = editor.state.doc.nodeAt(cellPos);
+        if (!cellNode) continue;
+        tr.setNodeMarkup(cellPos, targetType, cellNode.attrs);
+      }
+
+      editor.view.dispatch(tr);
+    } catch {
+      // table may have changed
+    }
+
+    setMenuOpen(false);
+  };
+
   const handleDeleteRow = () => {
     if (!currentRow) return;
     const { rowIndex, tableStart } = currentRow;
@@ -626,6 +675,12 @@ export function TableRowMenu({ editor }: TableRowMenuProps) {
           onMouseDown={e => e.preventDefault()}
           className={menu()}
         >
+          {isFirstRow && (
+            <MenuButton active={isHeaderRow()} onClick={handleToggleHeaderRow}>
+              <CrownSimpleIcon size={16} weight="bold" />
+              Header row
+            </MenuButton>
+          )}
           <MenuButton onClick={handleInsertAbove}>
             <PlusIcon size={16} weight="bold" />
             Insert row above