/**
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 *
 */

import type {TableCellNode} from './LexicalTableCellNode';
import type {
  TableDOMCell,
  TableDOMRows,
  TableObservers,
} from './LexicalTableObserver';
import type {
  TableMapType,
  TableMapValueType,
  TableSelection,
} from './LexicalTableSelection';
import type {
  BaseSelection,
  CaretDirection,
  ChildCaret,
  EditorState,
  ElementNode,
  LexicalCommand,
  LexicalEditor,
  LexicalNode,
  PointCaret,
  RangeSelection,
  SiblingCaret,
} from 'lexical';

import {
  $getClipboardDataFromSelection,
  copyToClipboard,
} from '@lexical/clipboard';
import invariant from '@lexical/internal/invariant';
import {
  $findMatchingParent,
  addClassNamesToElement,
  objectKlassEquals,
  removeClassNamesFromElement,
} from '@lexical/utils';
import {
  $caretFromPoint,
  $createParagraphNode,
  $createRangeSelectionFromDom,
  $createTextNode,
  $extendCaretToRange,
  $getAdjacentChildCaret,
  $getChildCaret,
  $getNearestNodeFromDOMNode,
  $getNodeByKeyOrThrow,
  $getPreviousSelection,
  $getSelection,
  $getSiblingCaret,
  $isChildCaret,
  $isElementNode,
  $isExtendableTextPointCaret,
  $isRangeSelection,
  $isRootOrShadowRoot,
  $isSiblingCaret,
  $isTextNode,
  $normalizeCaret,
  $setPointFromCaret,
  $setSelection,
  COMMAND_PRIORITY_HIGH,
  CONTROLLED_TEXT_INSERTION_COMMAND,
  CUT_COMMAND,
  DELETE_CHARACTER_COMMAND,
  DELETE_LINE_COMMAND,
  DELETE_WORD_COMMAND,
  FOCUS_COMMAND,
  FORMAT_ELEMENT_COMMAND,
  FORMAT_TEXT_COMMAND,
  getDOMSelection,
  INSERT_PARAGRAPH_COMMAND,
  IS_FIREFOX,
  isDOMNode,
  isHTMLElement,
  KEY_ARROW_DOWN_COMMAND,
  KEY_ARROW_LEFT_COMMAND,
  KEY_ARROW_RIGHT_COMMAND,
  KEY_ARROW_UP_COMMAND,
  KEY_BACKSPACE_COMMAND,
  KEY_DELETE_COMMAND,
  KEY_ESCAPE_COMMAND,
  KEY_TAB_COMMAND,
  SELECTION_CHANGE_COMMAND,
} from 'lexical';

import {$isTableCellNode} from './LexicalTableCellNode';
import {
  $getElementForTableNode,
  $isScrollableTablesActive,
  $isTableNode,
  TableNode,
} from './LexicalTableNode';
import {TableDOMTable, TableObserver} from './LexicalTableObserver';
import {$isTableRowNode} from './LexicalTableRowNode';
import {$isTableSelection} from './LexicalTableSelection';
import {
  $computeTableCellRectBoundary,
  $computeTableCellRectSpans,
  $computeTableMap,
  $getNodeTriplet,
  TableCellRectBoundary,
} from './LexicalTableUtils';

const LEXICAL_ELEMENT_KEY = '__lexicalTableSelection';

const isPointerDownOnEvent = (event: PointerEvent) => {
  return (event.buttons & 1) === 1;
};

export function isHTMLTableElement(el: unknown): el is HTMLTableElement {
  return isHTMLElement(el) && el.nodeName === 'TABLE';
}

export function getTableElement<T extends HTMLElement | null>(
  tableNode: TableNode,
  dom: T,
): HTMLTableElementWithWithTableSelectionState | (T & null) {
  if (!dom) {
    return dom as T & null;
  }
  const element: null | HTMLTableElementWithWithTableSelectionState =
    isHTMLTableElement(dom) ? dom : dom.querySelector('table');
  invariant(
    isHTMLTableElement(element),
    'getTableElement: Expecting table in DOM node for %s of type %s with key %s, not %s',
    tableNode.constructor.name,
    tableNode.getType(),
    tableNode.getKey(),
    dom.nodeName,
  );
  return element;
}

export function getEditorWindow(editor: LexicalEditor): Window | null {
  return editor._window;
}

export function $findParentTableCellNodeInTable(
  tableNode: LexicalNode,
  node: LexicalNode | null,
): TableCellNode | null {
  for (
    let currentNode = node, lastTableCellNode: TableCellNode | null = null;
    currentNode !== null;
    currentNode = currentNode.getParent()
  ) {
    if (tableNode.is(currentNode)) {
      return lastTableCellNode;
    } else if ($isTableCellNode(currentNode)) {
      lastTableCellNode = currentNode;
    }
  }
  return null;
}

const ARROW_KEY_COMMANDS_WITH_DIRECTION = [
  [KEY_ARROW_DOWN_COMMAND, 'down'],
  [KEY_ARROW_UP_COMMAND, 'up'],
  [KEY_ARROW_LEFT_COMMAND, 'backward'],
  [KEY_ARROW_RIGHT_COMMAND, 'forward'],
] as const;
const DELETE_TEXT_COMMANDS = [
  DELETE_WORD_COMMAND,
  DELETE_LINE_COMMAND,
  DELETE_CHARACTER_COMMAND,
] as const;
const DELETE_KEY_COMMANDS = [
  KEY_BACKSPACE_COMMAND,
  KEY_DELETE_COMMAND,
] as const;

export function registerTableWindowHandlers(
  editor: LexicalEditor,
  tableObservers: TableObservers,
) {
  // Use registerRootListener so the pointerdown handler is (re)attached
  // whenever the root element is set. This is required for the Extension API,
  // where register() runs before the ContentEditable mounts and getRootElement()
  // is still null.
  return editor.registerRootListener(rootElement => {
    if (rootElement === null) {
      return;
    }
    const editorWindow = editor._window;
    if (editorWindow === null) {
      return;
    }

    const pointerDownCallback = (event: PointerEvent) => {
      const target = event.target;
      if (
        event.button !== 0 ||
        !isDOMNode(target) ||
        !rootElement.contains(target)
      ) {
        return;
      }
      const selectionInfo = getTableObserverFromCellNode(target);

      editor.update(() => {
        // Clear highlights from all tables (even one we're actively clicking on)
        const selection = $getSelection();
        if ($isTableSelection(selection)) {
          for (const [observer] of tableObservers.observers.values()) {
            observer.$clearHighlight(false);
          }
          $setSelection(null);
          editor.dispatchCommand(SELECTION_CHANGE_COMMAND, undefined);
        }
        if (!selectionInfo) {
          return;
        }
        const {tableObserver, tableElement, cellElement} = selectionInfo;
        $handleTableClick(
          editor,
          event,
          cellElement,
          tableElement,
          tableObserver,
          tableObservers,
        );
      });
    };

    editorWindow.addEventListener('pointerdown', pointerDownCallback);
    return () => {
      editorWindow.removeEventListener('pointerdown', pointerDownCallback);
    };
  });
}

function $handleTableClick(
  editor: LexicalEditor,
  event: PointerEvent,
  selectedDOMCell: TableDOMCell,
  tableElement: HTMLTableElementWithWithTableSelectionState,
  tableObserver: TableObserver,
  tableObservers: TableObservers,
) {
  const editorWindow = editor._window;
  if (!editorWindow) {
    return;
  }
  const createPointerHandlers = (startingCell: TableDOMCell | null) => {
    if (tableObserver.isSelecting) {
      return;
    }
    tableObserver.isSelecting = true;

    // Set anchor immediately if starting cell provided (handles direct drag without click)
    if (startingCell !== null && tableObserver.anchorCell === null) {
      editor.update(() => {
        tableObserver.$setAnchorCellForSelection(startingCell);
      });
    }

    const onPointerUp = () => {
      tableObserver.isSelecting = false;
      editorWindow.removeEventListener('pointerup', onPointerUp);
      editorWindow.removeEventListener('pointermove', onPointerMove);
    };

    const onPointerMove = (moveEvent: PointerEvent) => {
      if (!isPointerDownOnEvent(moveEvent) && tableObserver.isSelecting) {
        tableObserver.isSelecting = false;
        editorWindow.removeEventListener('pointerup', onPointerUp);
        editorWindow.removeEventListener('pointermove', onPointerMove);
        return;
      }
      if (!isDOMNode(moveEvent.target)) {
        return;
      }
      let focusCell: null | TableDOMCell = null;
      // In firefox the moveEvent.target may be captured so we must always
      // consult the coordinates #7245
      const override = !(IS_FIREFOX || tableElement.contains(moveEvent.target));
      if (override) {
        focusCell = getDOMCellInTableFromTarget(tableElement, moveEvent.target);
      } else {
        for (const el of document.elementsFromPoint(
          moveEvent.clientX,
          moveEvent.clientY,
        )) {
          focusCell = getDOMCellInTableFromTarget(tableElement, el);
          if (focusCell) {
            break;
          }
        }
      }
      if (focusCell) {
        const anchorCell = focusCell;
        // Fallback: set anchor if still missing (handles race conditions)
        if (tableObserver.anchorCell === null) {
          editor.update(() => {
            tableObserver.$setAnchorCellForSelection(anchorCell);
          });
        }
        if (
          tableObserver.focusCell === null ||
          focusCell.elem !== tableObserver.focusCell.elem
        ) {
          tableObservers.setNextFocus({
            focusCell,
            override,
            tableKey: tableObserver.tableNodeKey,
          });
          editor.dispatchCommand(SELECTION_CHANGE_COMMAND, undefined);
        }
      }
    };
    editorWindow.addEventListener(
      'pointerup',
      onPointerUp,
      tableObserver.listenerOptions,
    );
    editorWindow.addEventListener(
      'pointermove',
      onPointerMove,
      tableObserver.listenerOptions,
    );
  };

  tableObserver.pointerType = event.pointerType;
  const tableNode = $getNodeByKeyOrThrow<TableNode>(tableObserver.tableNodeKey);
  const prevSelection = $getPreviousSelection();
  // We can't trust Firefox to do the right thing with the selection and
  // we don't have a proper state machine to do this "correctly" but
  // if we go ahead and make the table selection now it will work
  if (
    IS_FIREFOX &&
    event.shiftKey &&
    $isSelectionInTable(prevSelection, tableNode) &&
    ($isRangeSelection(prevSelection) || $isTableSelection(prevSelection))
  ) {
    const prevAnchorNode = prevSelection.anchor.getNode();
    const prevAnchorCell = $findParentTableCellNodeInTable(
      tableNode,
      prevSelection.anchor.getNode(),
    );
    if (prevAnchorCell) {
      tableObserver.$setAnchorCellForSelection(
        $getObserverCellFromCellNodeOrThrow(tableObserver, prevAnchorCell),
      );
      tableObserver.$setFocusCellForSelection(selectedDOMCell);
      stopEvent(event);
    } else {
      const newSelection = tableNode.isBefore(prevAnchorNode)
        ? tableNode.selectStart()
        : tableNode.selectEnd();
      newSelection.anchor.set(
        prevSelection.anchor.key,
        prevSelection.anchor.offset,
        prevSelection.anchor.type,
      );
    }
  } else {
    // Only set anchor cell for selection if this is not a simple touch tap
    // Touch taps should not initiate table selection mode
    if (event.pointerType !== 'touch') {
      tableObserver.$setAnchorCellForSelection(selectedDOMCell);
    }
  }

  // Pass the target cell to createPointerHandlers so it can be used as anchor
  // if user drags directly without clicking first
  createPointerHandlers(selectedDOMCell);
}

export function applyTableHandlers(
  tableNode: TableNode,
  element: HTMLElement,
  editor: LexicalEditor,
  hasTabHandler: boolean,
  tableObservers: TableObservers,
): TableObserver {
  const rootElement = editor.getRootElement();
  const editorWindow = getEditorWindow(editor);
  invariant(
    rootElement !== null && editorWindow !== null,
    'applyTableHandlers: editor has no root element set',
  );

  const tableObserver = new TableObserver(editor, tableNode.getKey());

  const tableElement = getTableElement(tableNode, element);
  attachTableObserverToTableElement(tableElement, tableObserver);
  tableObserver.listenersToRemove.add(() =>
    detachTableObserverFromTableElement(tableElement, tableObserver),
  );

  const onTripleClick = (event: MouseEvent) => {
    if (event.detail >= 3 && isDOMNode(event.target)) {
      const targetCell = getDOMCellFromTarget(event.target);
      if (targetCell !== null) {
        event.preventDefault();
      }
    }
  };
  tableElement.addEventListener(
    'mousedown',
    onTripleClick,
    tableObserver.listenerOptions,
  );
  tableObserver.listenersToRemove.add(() => {
    tableElement.removeEventListener('mousedown', onTripleClick);
  });

  for (const [command, direction] of ARROW_KEY_COMMANDS_WITH_DIRECTION) {
    tableObserver.listenersToRemove.add(
      editor.registerCommand(
        command,
        event =>
          $handleArrowKey(
            editor,
            event,
            direction,
            tableNode,
            tableObserver,
            tableObservers,
          ),
        COMMAND_PRIORITY_HIGH,
      ),
    );
  }

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      KEY_ESCAPE_COMMAND,
      event => {
        const selection = $getSelection();
        if ($isTableSelection(selection)) {
          const focusCellNode = $findParentTableCellNodeInTable(
            tableNode,
            selection.focus.getNode(),
          );
          if (focusCellNode !== null) {
            stopEvent(event);
            focusCellNode.selectEnd();
            return true;
          }
        }

        return false;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  const deleteTextHandler = (command: LexicalCommand<boolean>) => () => {
    const selection = $getSelection();

    if (!$isSelectionInTable(selection, tableNode)) {
      return false;
    }

    if ($isTableSelection(selection)) {
      tableObserver.$clearText();

      return true;
    } else if ($isRangeSelection(selection)) {
      const tableCellNode = $findParentTableCellNodeInTable(
        tableNode,
        selection.anchor.getNode(),
      );

      if (!$isTableCellNode(tableCellNode)) {
        return false;
      }

      const anchorNode = selection.anchor.getNode();
      const focusNode = selection.focus.getNode();
      const isAnchorInside = tableNode.isParentOf(anchorNode);
      const isFocusInside = tableNode.isParentOf(focusNode);

      const selectionContainsPartialTable =
        (isAnchorInside && !isFocusInside) ||
        (isFocusInside && !isAnchorInside);

      if (selectionContainsPartialTable) {
        tableObserver.$clearText();
        return true;
      }

      const nearestElementNode = $findMatchingParent(
        selection.anchor.getNode(),
        n => $isElementNode(n),
      );

      const topLevelCellElementNode =
        nearestElementNode &&
        $findMatchingParent(
          nearestElementNode,
          n => $isElementNode(n) && $isTableCellNode(n.getParent()),
        );

      if (
        !$isElementNode(topLevelCellElementNode) ||
        !$isElementNode(nearestElementNode)
      ) {
        return false;
      }

      if (
        command === DELETE_LINE_COMMAND &&
        topLevelCellElementNode.getPreviousSibling() === null
      ) {
        // TODO: Fix Delete Line in Table Cells.
        return true;
      }
    }

    return false;
  };

  for (const command of DELETE_TEXT_COMMANDS) {
    tableObserver.listenersToRemove.add(
      editor.registerCommand(
        command,
        deleteTextHandler(command),
        COMMAND_PRIORITY_HIGH,
      ),
    );
  }

  const $deleteCellHandler = (
    event: KeyboardEvent | ClipboardEvent | null,
  ): boolean => {
    const selection = $getSelection();
    if (!($isTableSelection(selection) || $isRangeSelection(selection))) {
      return false;
    }

    // If the selection is inside the table but should remove the whole table
    // we expand the selection so that both the anchor and focus are outside
    // the table and the editor's command listener will handle the delete
    const isAnchorInside = tableNode.isParentOf(selection.anchor.getNode());
    const isFocusInside = tableNode.isParentOf(selection.focus.getNode());
    if (isAnchorInside !== isFocusInside) {
      const tablePoint = isAnchorInside ? 'anchor' : 'focus';
      const outerPoint = isAnchorInside ? 'focus' : 'anchor';
      // Preserve the outer point
      const {key, offset, type} = selection[outerPoint];
      // Expand the selection around the table
      const newSelection =
        tableNode[
          selection[tablePoint].isBefore(selection[outerPoint])
            ? 'selectPrevious'
            : 'selectNext'
        ]();
      // Restore the outer point of the selection
      newSelection[outerPoint].set(key, offset, type);
      // Let the base implementation handle the rest
      return false;
    }

    if (!$isSelectionInTable(selection, tableNode)) {
      return false;
    }

    if ($isTableSelection(selection)) {
      if (event) {
        event.preventDefault();
        event.stopPropagation();
      }
      tableObserver.$clearText();

      return true;
    }

    return false;
  };

  for (const command of DELETE_KEY_COMMANDS) {
    tableObserver.listenersToRemove.add(
      editor.registerCommand(
        command,
        $deleteCellHandler,
        COMMAND_PRIORITY_HIGH,
      ),
    );
  }

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      CUT_COMMAND,
      event => {
        const selection = $getSelection();
        if (selection) {
          if (!($isTableSelection(selection) || $isRangeSelection(selection))) {
            return false;
          }
          // Copying to the clipboard is async so we must capture the data
          // before we delete it
          void copyToClipboard(
            editor,
            objectKlassEquals(event, ClipboardEvent) ? event : null,
            $getClipboardDataFromSelection(selection),
          );
          const intercepted = $deleteCellHandler(event);
          if ($isRangeSelection(selection)) {
            selection.removeText();
            return true;
          }
          return intercepted;
        }
        return false;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      FORMAT_TEXT_COMMAND,
      payload => {
        const selection = $getSelection();

        if (!$isSelectionInTable(selection, tableNode)) {
          return false;
        }

        if ($isTableSelection(selection)) {
          tableObserver.$formatCells(payload);

          return true;
        } else if ($isRangeSelection(selection)) {
          const tableCellNode = $findMatchingParent(
            selection.anchor.getNode(),
            n => $isTableCellNode(n),
          );

          if (!$isTableCellNode(tableCellNode)) {
            return false;
          }
        }

        return false;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      FORMAT_ELEMENT_COMMAND,
      formatType => {
        const selection = $getSelection();
        if (
          !$isTableSelection(selection) ||
          !$isSelectionInTable(selection, tableNode)
        ) {
          return false;
        }

        const anchorNode = selection.anchor.getNode();
        const focusNode = selection.focus.getNode();
        if (!$isTableCellNode(anchorNode) || !$isTableCellNode(focusNode)) {
          return false;
        }

        // Align the table if the entire table is selected
        if ($isFullTableSelection(selection, tableNode)) {
          tableNode.setFormat(formatType);
          return true;
        }

        const [tableMap, anchorCell, focusCell] = $computeTableMap(
          tableNode,
          anchorNode,
          focusNode,
        );
        const maxRow = Math.max(
          anchorCell.startRow + anchorCell.cell.__rowSpan - 1,
          focusCell.startRow + focusCell.cell.__rowSpan - 1,
        );
        const maxColumn = Math.max(
          anchorCell.startColumn + anchorCell.cell.__colSpan - 1,
          focusCell.startColumn + focusCell.cell.__colSpan - 1,
        );
        const minRow = Math.min(anchorCell.startRow, focusCell.startRow);
        const minColumn = Math.min(
          anchorCell.startColumn,
          focusCell.startColumn,
        );
        const visited = new Set<TableCellNode>();
        for (let i = minRow; i <= maxRow; i++) {
          for (let j = minColumn; j <= maxColumn; j++) {
            const cell = tableMap[i][j].cell;
            if (visited.has(cell)) {
              continue;
            }
            visited.add(cell);
            cell.setFormat(formatType);

            const cellChildren = cell.getChildren();
            for (let k = 0; k < cellChildren.length; k++) {
              const child = cellChildren[k];
              if ($isElementNode(child) && !child.isInline()) {
                child.setFormat(formatType);
              }
            }
          }
        }
        return true;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      CONTROLLED_TEXT_INSERTION_COMMAND,
      payload => {
        const selection = $getSelection();

        if (!$isSelectionInTable(selection, tableNode)) {
          return false;
        }

        if ($isTableSelection(selection)) {
          tableObserver.$clearHighlight();

          return false;
        } else if ($isRangeSelection(selection)) {
          const tableCellNode = $findMatchingParent(
            selection.anchor.getNode(),
            n => $isTableCellNode(n),
          );

          if (!$isTableCellNode(tableCellNode)) {
            return false;
          }

          if (typeof payload === 'string') {
            const edgePosition = $getTableEdgeCursorPosition(
              editor,
              selection,
              tableNode,
            );
            if (edgePosition) {
              $insertParagraphAtTableEdge(edgePosition, tableNode, [
                $createTextNode(payload),
              ]);
              return true;
            }
          }
        }

        return false;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  if (hasTabHandler) {
    tableObserver.listenersToRemove.add(
      editor.registerCommand(
        KEY_TAB_COMMAND,
        event => {
          const selection = $getSelection();
          if (
            !$isRangeSelection(selection) ||
            !selection.isCollapsed() ||
            !$isSelectionInTable(selection, tableNode)
          ) {
            return false;
          }

          const tableCellNode = $findCellNode(selection.anchor.getNode());
          if (
            tableCellNode === null ||
            !tableNode.is($findTableNode(tableCellNode))
          ) {
            return false;
          }

          stopEvent(event);
          $selectAdjacentCell(
            tableCellNode,
            event.shiftKey ? 'previous' : 'next',
          );

          return true;
        },
        COMMAND_PRIORITY_HIGH,
      ),
    );
  }

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      FOCUS_COMMAND,
      payload => {
        return tableNode.isSelected();
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  tableObserver.listenersToRemove.add(
    editor.registerCommand(
      INSERT_PARAGRAPH_COMMAND,
      () => {
        const selection = $getSelection();
        if (
          !$isRangeSelection(selection) ||
          !selection.isCollapsed() ||
          !$isSelectionInTable(selection, tableNode)
        ) {
          return false;
        }
        const edgePosition = $getTableEdgeCursorPosition(
          editor,
          selection,
          tableNode,
        );
        if (edgePosition) {
          $insertParagraphAtTableEdge(edgePosition, tableNode);
          return true;
        }
        return false;
      },
      COMMAND_PRIORITY_HIGH,
    ),
  );

  return tableObserver;
}

/** @internal */
export function $handleTableSelectionChangeCommand(
  tableObservers: TableObservers,
  editor: LexicalEditor,
) {
  const selection = $getSelection();
  const prevSelection = $getPreviousSelection();

  const nextFocus = tableObservers.getAndClearNextFocus();
  if (nextFocus !== null) {
    const {tableKey, focusCell} = nextFocus;
    const observerAndTable = tableObservers.observers.get(tableKey);
    invariant(
      !!observerAndTable,
      'tableObserver not found for tableKey: %s',
      tableKey,
    );
    const [tableObserver] = observerAndTable;
    if (
      $isTableSelection(selection) &&
      selection.tableKey === tableObserver.tableNodeKey
    ) {
      if (
        focusCell.x === tableObserver.focusX &&
        focusCell.y === tableObserver.focusY
      ) {
        // The selection is already the correct table selection
        return false;
      } else {
        tableObserver.$setFocusCellForSelection(focusCell);
        return true;
      }
    } else if (
      tableObserver.anchorCell !== null &&
      tableObserver.anchorCellNodeKey !== null &&
      focusCell.elem !== tableObserver.anchorCell.elem &&
      tableObserver.tableSelection !== null
    ) {
      // The selection has crossed cells
      // If we have an anchor cell set and tableSelection initialized,
      // we have all the necessary state to create the selection.
      // The presence of nextFocus means we're dragging, so process it.
      // Use ignoreStart=true to ensure isHighlightingCells is set correctly
      // on the first drag attempt, especially when switching columns.
      tableObserver.$setFocusCellForSelection(focusCell, true);
      return true;
    }
  }
  const shouldCheckSelectionForTable =
    tableObservers.getAndClearShouldCheckSelectionForTable();
  // If they pressed the down arrow with the selection outside of the
  // table, and then the selection ends up in the table but not in the
  // first cell, then move the selection to the first cell.
  if (
    !!shouldCheckSelectionForTable &&
    $isRangeSelection(prevSelection) &&
    $isRangeSelection(selection) &&
    selection.isCollapsed()
  ) {
    const tableNode = $getNodeByKeyOrThrow<TableNode>(
      shouldCheckSelectionForTable,
    );
    const anchor = selection.anchor.getNode();
    const firstRow = tableNode.getFirstChild();
    const anchorCell = $findCellNode(anchor);
    if (anchorCell !== null && $isTableRowNode(firstRow)) {
      const firstCell = firstRow.getFirstChild();
      if (
        $isTableCellNode(firstCell) &&
        tableNode.is(
          $findMatchingParent(
            anchorCell,
            node => node.is(tableNode) || node.is(firstCell),
          ),
        )
      ) {
        // The selection moved to the table, but not in the first cell
        firstCell.selectStart();
        return true;
      }
    }
  }

  if ($isTableSelection(selection)) {
    $fixTableSelectionForSelectedTable(editor, selection);
  }

  if ($isRangeSelection(selection)) {
    $fixRangeSelectionForSelectedTable(selection, tableObservers);
  }

  // Generic selection logic that runs across every table observer when the selection changes.
  // Note: the selection might have changed in the code above, which re-dispatches the selection change command
  // and gets handled here on the second pass. This should be refactored.
  const tableNodesAndObservers = Array.from(
    tableObservers.observers.entries(),
  ).map(([tableKey, [tableObserver]]) => ({
    tableNode: $getNodeByKeyOrThrow<TableNode>(tableKey),
    tableObserver,
  }));
  for (const {tableNode, tableObserver} of tableNodesAndObservers) {
    $syncTableSelectionState(editor, tableNode, tableObserver);
  }

  return false;
}

/**
 * Handles cases where range selections cross into, out of, or within tables.
 */
function $fixRangeSelectionForSelectedTable(
  selection: RangeSelection,
  tableObservers: TableObservers,
) {
  const prevSelection = $getPreviousSelection();
  const {anchor, focus} = selection;
  const anchorNode = anchor.getNode();
  const focusNode = focus.getNode();
  // Using explicit comparison with table node to ensure it's not a nested table
  // as in that case we'll leave selection resolving to that table
  const anchorCellNode = $findCellNode(anchorNode);
  const focusCellNode = $findCellNode(focusNode);
  const anchorCellTable = anchorCellNode
    ? $findTableNode(anchorCellNode)
    : null;
  const focusCellTable = focusCellNode ? $findTableNode(focusCellNode) : null;
  const isBackward = selection.isBackward();

  const isSameTable =
    anchorCellNode &&
    focusCellNode &&
    anchorCellTable &&
    focusCellTable &&
    anchorCellTable.is(focusCellTable);

  // The focus should be moved (to cover the whole focus table) if it is moved outside of the anchor's table.
  // For example, when dragging from outside a table into it.
  const shouldMoveFocus =
    focusCellTable &&
    (!anchorCellTable || anchorCellTable.isParentOf(focusCellTable));
  // The anchor should be moved (to cover the whole anchor table) if the focus is moved outside of the anchor table.
  // For example, when dragging from inside a table out of it.
  const shouldMoveAnchor =
    anchorCellTable &&
    (!focusCellTable || focusCellTable.isParentOf(anchorCellTable));

  if (shouldMoveFocus) {
    // Select the whole focus table.
    const newSelection = selection.clone();
    const [tableMap] = $computeTableMap(
      focusCellTable,
      focusCellNode!,
      focusCellNode!,
    );
    const firstCell = tableMap[0][0].cell;
    const lastCell = tableMap[tableMap.length - 1].at(-1)!.cell;
    newSelection.focus.set(
      isBackward ? firstCell.getKey() : lastCell.getKey(),
      isBackward ? 0 : lastCell.getChildrenSize(),
      'element',
    );
    $setSelection(newSelection);
  } else if (shouldMoveAnchor) {
    // Select the whole anchor table.
    const newSelection = selection.clone();
    const [tableMap] = $computeTableMap(
      anchorCellTable,
      anchorCellNode!,
      anchorCellNode!,
    );
    const firstCell = tableMap[0][0].cell;
    const lastCell = tableMap[tableMap.length - 1].at(-1)!.cell;
    newSelection.anchor.set(
      isBackward ? lastCell.getKey() : firstCell.getKey(),
      isBackward ? lastCell.getChildrenSize() : 0,
      'element',
    );
    $setSelection(newSelection);
  } else if (isSameTable) {
    // Handle case when selection spans across multiple cells but still
    // has range selection, then we convert it into table selection
    // For example, this fires when dragging up from first cell, outside of the table, or when clicking a cell
    // then shift-clicking another cell.
    const observerInfo = tableObservers.observers.get(anchorCellTable.getKey());
    invariant(
      !!observerInfo,
      'tableObserver not found for tableKey: %s',
      anchorCellTable.getKey(),
    );
    const [tableObserver] = observerInfo;
    if (!anchorCellNode.is(focusCellNode)) {
      tableObserver.$setAnchorCellForSelection(
        $getObserverCellFromCellNodeOrThrow(tableObserver, anchorCellNode),
      );
      tableObserver.$setFocusCellForSelection(
        $getObserverCellFromCellNodeOrThrow(tableObserver, focusCellNode),
        true,
      );
    }

    // Handle case when the pointer type is touch and the current and
    // previous selection are collapsed, and the previous anchor and current
    // focus cell nodes are different, then we convert it into table selection
    // However, only do this if the table observer is actively selecting (user dragging)
    // to prevent unwanted selections when simply tapping between cells on mobile
    if (
      tableObserver.pointerType === 'touch' &&
      tableObserver.isSelecting &&
      selection.isCollapsed() &&
      $isRangeSelection(prevSelection) &&
      prevSelection.isCollapsed()
    ) {
      const prevAnchorCellNode = $findCellNode(prevSelection.anchor.getNode());
      if (prevAnchorCellNode && !prevAnchorCellNode.is(focusCellNode)) {
        tableObserver.$setAnchorCellForSelection(
          $getObserverCellFromCellNodeOrThrow(
            tableObserver,
            prevAnchorCellNode,
          ),
        );
        tableObserver.$setFocusCellForSelection(
          $getObserverCellFromCellNodeOrThrow(tableObserver, focusCellNode),
          true,
        );
        tableObserver.pointerType = null;
      }
    }
  }
}

/**
 * Ensures that a TableSelection is automatically changed to a RangeSelection when the selection goes outside of the table.
 */
function $fixTableSelectionForSelectedTable(
  editor: LexicalEditor,
  selection: TableSelection,
) {
  const editorWindow = getEditorWindow(editor);
  const prevSelection = $getPreviousSelection();
  if (!selection.is(prevSelection)) {
    return;
  }
  const tableNode = $getNodeByKeyOrThrow<TableNode>(selection.tableKey);
  // if selection goes outside of the table we need to change it to Range selection
  const domSelection = getDOMSelection(editorWindow);
  if (domSelection && domSelection.anchorNode && domSelection.focusNode) {
    const focusNode = $getNearestNodeFromDOMNode(domSelection.focusNode);
    const isFocusOutside = focusNode && !tableNode.isParentOf(focusNode);

    const anchorNode = $getNearestNodeFromDOMNode(domSelection.anchorNode);
    const isAnchorInside = anchorNode && tableNode.isParentOf(anchorNode);

    if (isFocusOutside && isAnchorInside && domSelection.rangeCount > 0) {
      const newSelection = $createRangeSelectionFromDom(domSelection, editor);
      if (newSelection) {
        newSelection.anchor.set(
          tableNode.getKey(),
          selection.isBackward() ? tableNode.getChildrenSize() : 0,
          'element',
        );
        domSelection.removeAllRanges();
        $setSelection(newSelection);
      }
    }
  }
}

// Handle keeping the table observer/DOM in sync with the selection.
function $syncTableSelectionState(
  editor: LexicalEditor,
  tableNode: TableNode,
  tableObserver: TableObserver,
) {
  const selection = $getSelection();
  const prevSelection = $getPreviousSelection();
  if (
    selection &&
    !selection.is(prevSelection) &&
    ($isTableSelection(selection) || $isTableSelection(prevSelection)) &&
    tableObserver.tableSelection &&
    !tableObserver.tableSelection.is(prevSelection)
  ) {
    if (
      $isTableSelection(selection) &&
      selection.tableKey === tableObserver.tableNodeKey
    ) {
      tableObserver.$updateTableTableSelection(selection);
    } else if (
      !$isTableSelection(selection) &&
      $isTableSelection(prevSelection) &&
      prevSelection.tableKey === tableObserver.tableNodeKey
    ) {
      tableObserver.$updateTableTableSelection(null);
    }
  }

  if (tableObserver.hasHijackedSelectionStyles && !tableNode.isSelected()) {
    $removeHighlightStyleToTable(editor, tableObserver);
  } else if (
    !tableObserver.hasHijackedSelectionStyles &&
    tableNode.isSelected()
  ) {
    $addHighlightStyleToTable(editor, tableObserver);
  }
}

export type HTMLTableElementWithWithTableSelectionState = HTMLTableElement & {
  [LEXICAL_ELEMENT_KEY]?: TableObserver | undefined;
};

export function detachTableObserverFromTableElement(
  tableElement: HTMLTableElementWithWithTableSelectionState,
  tableObserver: TableObserver,
) {
  if (getTableObserverFromTableElement(tableElement) === tableObserver) {
    delete tableElement[LEXICAL_ELEMENT_KEY];
  }
}

export function attachTableObserverToTableElement(
  tableElement: HTMLTableElementWithWithTableSelectionState,
  tableObserver: TableObserver,
) {
  invariant(
    getTableObserverFromTableElement(tableElement) === null,
    'tableElement already has an attached TableObserver',
  );
  tableElement[LEXICAL_ELEMENT_KEY] = tableObserver;
}

export function getTableObserverFromTableElement(
  tableElement: HTMLTableElementWithWithTableSelectionState,
): TableObserver | null {
  return tableElement[LEXICAL_ELEMENT_KEY] || null;
}

function getTableObserverFromCellNode(node: null | Node): {
  tableObserver: TableObserver;
  tableElement: HTMLTableElementWithWithTableSelectionState;
  cellElement: TableDOMCell;
} | null {
  const cellNode = getDOMCellFromTarget(node);
  if (cellNode === null) {
    return null;
  }
  let currentNode: ParentNode | Node | null = cellNode.elem;
  while (currentNode != null) {
    const nodeName = currentNode.nodeName;
    if (
      nodeName === 'TABLE' &&
      LEXICAL_ELEMENT_KEY in currentNode &&
      !!currentNode[LEXICAL_ELEMENT_KEY]
    ) {
      return {
        cellElement: cellNode,
        tableElement:
          currentNode as HTMLTableElementWithWithTableSelectionState,
        tableObserver: currentNode[LEXICAL_ELEMENT_KEY] as TableObserver,
      };
    }
    currentNode = currentNode.parentNode;
  }
  return null;
}

export function getDOMCellFromTarget(node: null | Node): TableDOMCell | null {
  let currentNode: ParentNode | Node | null = node;

  while (currentNode != null) {
    const nodeName = currentNode.nodeName;

    if (nodeName === 'TD' || nodeName === 'TH') {
      // @ts-expect-error: internal field
      const cell = currentNode._cell;

      if (cell === undefined) {
        return null;
      }

      return cell;
    }

    currentNode = currentNode.parentNode;
  }

  return null;
}

export function getDOMCellInTableFromTarget(
  table: HTMLTableElementWithWithTableSelectionState,
  node: null | Node,
): TableDOMCell | null {
  if (!table.contains(node)) {
    return null;
  }
  let cell: null | TableDOMCell = null;
  for (
    let currentNode: ParentNode | Node | null = node;
    currentNode != null;
    currentNode = currentNode.parentNode
  ) {
    if (currentNode === table) {
      return cell;
    }
    const nodeName = currentNode.nodeName;
    if (nodeName === 'TD' || nodeName === 'TH') {
      // @ts-expect-error: internal field
      cell = currentNode._cell || null;
    }
  }
  return null;
}

export function doesTargetContainText(node: Node): boolean {
  const currentNode: ParentNode | Node | null = node;

  if (currentNode !== null) {
    const nodeName = currentNode.nodeName;

    if (nodeName === 'SPAN') {
      return true;
    }
  }
  return false;
}

export function getTable(
  tableNode: TableNode,
  dom: HTMLElement,
): TableDOMTable {
  const tableElement = getTableElement(tableNode, dom);
  const domRows: TableDOMRows = [];
  const grid = {
    columns: 0,
    domRows,
    rows: 0,
  };
  let currentNode = tableElement.querySelector('tr') as ChildNode | null;
  let x = 0;
  let y = 0;
  domRows.length = 0;

  while (currentNode != null) {
    const nodeMame = currentNode.nodeName;

    if (nodeMame === 'TD' || nodeMame === 'TH') {
      const elem = currentNode as HTMLElement;
      const cell = {
        elem,
        hasBackgroundColor: elem.style.backgroundColor !== '',
        highlighted: false,
        x,
        y,
      };

      // @ts-expect-error: internal field
      currentNode._cell = cell;

      let row = domRows[y];
      if (row === undefined) {
        row = domRows[y] = [];
      }

      row[x] = cell;
    } else {
      const child = currentNode.firstChild;

      if (child != null) {
        currentNode = child;
        continue;
      }
    }

    const sibling = currentNode.nextSibling;

    if (sibling != null) {
      x++;
      currentNode = sibling;
      continue;
    }

    const parent = currentNode.parentNode;

    if (parent != null) {
      const parentSibling = parent.nextSibling;

      if (parentSibling == null) {
        break;
      }

      y++;
      x = 0;
      currentNode = parentSibling;
    }
  }

  grid.columns = x + 1;
  grid.rows = y + 1;

  return grid;
}

export function $updateDOMForSelection(
  editor: LexicalEditor,
  table: TableDOMTable,
  selection: TableSelection | RangeSelection | null,
) {
  const selectedCellNodes = new Set(selection ? selection.getNodes() : []);
  $forEachTableCell(table, (cell, lexicalNode) => {
    const elem = cell.elem;

    if (selectedCellNodes.has(lexicalNode)) {
      cell.highlighted = true;
      $addHighlightToDOM(editor, cell);
    } else {
      cell.highlighted = false;
      $removeHighlightFromDOM(editor, cell);
      if (!elem.getAttribute('style')) {
        elem.removeAttribute('style');
      }
    }
  });
}

export function $forEachTableCell(
  grid: TableDOMTable,
  cb: (
    cell: TableDOMCell,
    lexicalNode: LexicalNode,
    cords: {
      x: number;
      y: number;
    },
  ) => void,
) {
  const {domRows} = grid;

  for (let y = 0; y < domRows.length; y++) {
    const row = domRows[y];
    if (!row) {
      continue;
    }

    for (let x = 0; x < row.length; x++) {
      const cell = row[x];
      if (!cell) {
        continue;
      }
      const lexicalNode = $getNearestNodeFromDOMNode(cell.elem);

      if (lexicalNode !== null) {
        cb(cell, lexicalNode, {
          x,
          y,
        });
      }
    }
  }
}

export function $addHighlightStyleToTable(
  editor: LexicalEditor,
  tableSelection: TableObserver,
) {
  tableSelection.$disableHighlightStyle();
  $forEachTableCell(tableSelection.table, cell => {
    cell.highlighted = true;
    $addHighlightToDOM(editor, cell);
  });
}

export function $removeHighlightStyleToTable(
  editor: LexicalEditor,
  tableObserver: TableObserver,
) {
  tableObserver.$enableHighlightStyle();
  $forEachTableCell(tableObserver.table, cell => {
    const elem = cell.elem;
    cell.highlighted = false;
    $removeHighlightFromDOM(editor, cell);

    if (!elem.getAttribute('style')) {
      elem.removeAttribute('style');
    }
  });
}

function $selectAdjacentCell(
  tableCellNode: TableCellNode,
  direction: 'next' | 'previous',
) {
  const siblingMethod =
    direction === 'next' ? 'getNextSibling' : 'getPreviousSibling';
  const childMethod = direction === 'next' ? 'getFirstChild' : 'getLastChild';
  const sibling = tableCellNode[siblingMethod]();
  if ($isElementNode(sibling)) {
    return sibling.selectEnd();
  }
  const parentRow = $findMatchingParent(tableCellNode, $isTableRowNode);
  invariant(parentRow !== null, 'selectAdjacentCell: Cell not in table row');
  for (
    let nextRow = parentRow[siblingMethod]();
    $isTableRowNode(nextRow);
    nextRow = nextRow[siblingMethod]()
  ) {
    const child = nextRow[childMethod]();
    if ($isElementNode(child)) {
      return child.selectEnd();
    }
  }
  const parentTable = $findMatchingParent(parentRow, $isTableNode);
  invariant(parentTable !== null, 'selectAdjacentCell: Row not in table');
  return direction === 'next'
    ? parentTable.selectNext()
    : parentTable.selectPrevious();
}

type Direction = 'backward' | 'forward' | 'up' | 'down';

const selectTableNodeInDirection = (
  tableObserver: TableObserver,
  tableNode: TableNode,
  x: number,
  y: number,
  direction: Direction,
): boolean => {
  const isForward = direction === 'forward';

  switch (direction) {
    case 'backward':
    case 'forward':
      if (x !== (isForward ? tableObserver.table.columns - 1 : 0)) {
        selectTableCellNode(
          tableNode.getCellNodeFromCordsOrThrow(
            x + (isForward ? 1 : -1),
            y,
            tableObserver.table,
          ),
          isForward,
        );
      } else {
        if (y !== (isForward ? tableObserver.table.rows - 1 : 0)) {
          selectTableCellNode(
            tableNode.getCellNodeFromCordsOrThrow(
              isForward ? 0 : tableObserver.table.columns - 1,
              y + (isForward ? 1 : -1),
              tableObserver.table,
            ),
            isForward,
          );
        } else if (!isForward) {
          tableNode.selectPrevious();
        } else {
          tableNode.selectNext();
        }
      }

      return true;

    case 'up':
      if (y !== 0) {
        selectTableCellNode(
          tableNode.getCellNodeFromCordsOrThrow(x, y - 1, tableObserver.table),
          false,
        );
      } else {
        tableNode.selectPrevious();
      }

      return true;

    case 'down':
      if (y !== tableObserver.table.rows - 1) {
        selectTableCellNode(
          tableNode.getCellNodeFromCordsOrThrow(x, y + 1, tableObserver.table),
          true,
        );
      } else {
        tableNode.selectNext();
      }

      return true;
    default:
      return false;
  }
};

type Corner = ['minColumn' | 'maxColumn', 'minRow' | 'maxRow'];
function getCorner(
  rect: TableCellRectBoundary,
  cellValue: TableMapValueType,
): Corner | null {
  let colName: 'minColumn' | 'maxColumn';
  let rowName: 'minRow' | 'maxRow';
  if (cellValue.startColumn === rect.minColumn) {
    colName = 'minColumn';
  } else if (
    cellValue.startColumn + cellValue.cell.__colSpan - 1 ===
    rect.maxColumn
  ) {
    colName = 'maxColumn';
  } else {
    return null;
  }
  if (cellValue.startRow === rect.minRow) {
    rowName = 'minRow';
  } else if (
    cellValue.startRow + cellValue.cell.__rowSpan - 1 ===
    rect.maxRow
  ) {
    rowName = 'maxRow';
  } else {
    return null;
  }
  return [colName, rowName];
}

function getCornerOrThrow(
  rect: TableCellRectBoundary,
  cellValue: TableMapValueType,
): Corner {
  const corner = getCorner(rect, cellValue);
  invariant(
    corner !== null,
    'getCornerOrThrow: cell %s is not at a corner of rect',
    cellValue.cell.getKey(),
  );
  return corner;
}

function oppositeCorner([colName, rowName]: Corner): Corner {
  return [
    colName === 'minColumn' ? 'maxColumn' : 'minColumn',
    rowName === 'minRow' ? 'maxRow' : 'minRow',
  ];
}

function cellAtCornerOrThrow(
  tableMap: TableMapType,
  rect: TableCellRectBoundary,
  [colName, rowName]: Corner,
): TableMapValueType {
  const rowNum = rect[rowName];
  const rowMap = tableMap[rowNum];
  invariant(
    rowMap !== undefined,
    'cellAtCornerOrThrow: %s = %s missing in tableMap',
    rowName,
    String(rowNum),
  );
  const colNum = rect[colName];
  const cell = rowMap[colNum];
  invariant(
    cell !== undefined,
    'cellAtCornerOrThrow: %s = %s missing in tableMap',
    colName,
    String(colNum),
  );
  return cell;
}

function $extractRectCorners(
  tableMap: TableMapType,
  anchorCellValue: TableMapValueType,
  newFocusCellValue: TableMapValueType,
) {
  // We are sure that the focus now either contracts or expands the rect
  // but both the anchor and focus might be moved to ensure a rectangle
  // given a potentially ragged merge shape
  const rect = $computeTableCellRectBoundary(
    tableMap,
    anchorCellValue,
    newFocusCellValue,
  );
  const anchorCorner = getCorner(rect, anchorCellValue);
  if (anchorCorner) {
    return [
      cellAtCornerOrThrow(tableMap, rect, anchorCorner),
      cellAtCornerOrThrow(tableMap, rect, oppositeCorner(anchorCorner)),
    ];
  }
  const newFocusCorner = getCorner(rect, newFocusCellValue);
  if (newFocusCorner) {
    return [
      cellAtCornerOrThrow(tableMap, rect, oppositeCorner(newFocusCorner)),
      cellAtCornerOrThrow(tableMap, rect, newFocusCorner),
    ];
  }
  // TODO this doesn't have to be arbitrary, use the closest corner instead
  const newAnchorCorner: Corner = ['minColumn', 'minRow'];
  return [
    cellAtCornerOrThrow(tableMap, rect, newAnchorCorner),
    cellAtCornerOrThrow(tableMap, rect, oppositeCorner(newAnchorCorner)),
  ];
}

function $adjustFocusInDirection(
  tableObserver: TableObserver,
  tableMap: TableMapType,
  anchorCellValue: TableMapValueType,
  focusCellValue: TableMapValueType,
  direction: Direction,
): boolean {
  const rect = $computeTableCellRectBoundary(
    tableMap,
    anchorCellValue,
    focusCellValue,
  );
  const spans = $computeTableCellRectSpans(tableMap, rect);
  const {topSpan, leftSpan, bottomSpan, rightSpan} = spans;
  const anchorCorner = getCornerOrThrow(rect, anchorCellValue);
  const [focusColumn, focusRow] = oppositeCorner(anchorCorner);
  let fCol = rect[focusColumn];
  let fRow = rect[focusRow];
  if (direction === 'forward') {
    fCol += focusColumn === 'maxColumn' ? 1 : leftSpan;
  } else if (direction === 'backward') {
    fCol -= focusColumn === 'minColumn' ? 1 : rightSpan;
  } else if (direction === 'down') {
    fRow += focusRow === 'maxRow' ? 1 : topSpan;
  } else if (direction === 'up') {
    fRow -= focusRow === 'minRow' ? 1 : bottomSpan;
  }
  const targetRowMap = tableMap[fRow];
  if (targetRowMap === undefined) {
    return false;
  }
  const newFocusCellValue = targetRowMap[fCol];
  if (newFocusCellValue === undefined) {
    return false;
  }
  // We can be certain that anchorCellValue and newFocusCellValue are
  // contained within the desired selection, but we are not certain if
  // they need to be expanded or not to maintain a rectangular shape
  const [finalAnchorCell, finalFocusCell] = $extractRectCorners(
    tableMap,
    anchorCellValue,
    newFocusCellValue,
  );
  const anchorDOM = $getObserverCellFromCellNodeOrThrow(
    tableObserver,
    finalAnchorCell.cell,
  )!;
  const focusDOM = $getObserverCellFromCellNodeOrThrow(
    tableObserver,
    finalFocusCell.cell,
  );
  tableObserver.$setAnchorCellForSelection(anchorDOM);
  tableObserver.$setFocusCellForSelection(focusDOM, true);
  return true;
}

function $isSelectionInTable(
  selection: null | BaseSelection,
  tableNode: TableNode,
): boolean {
  if ($isRangeSelection(selection) || $isTableSelection(selection)) {
    // TODO this should probably return false if there's an unrelated
    //      shadow root between the node and the table (e.g. another table,
    //      collapsible, etc.)
    const isAnchorInside = tableNode.isParentOf(selection.anchor.getNode());
    const isFocusInside = tableNode.isParentOf(selection.focus.getNode());

    return isAnchorInside && isFocusInside;
  }

  return false;
}

function $isFullTableSelection(
  selection: null | BaseSelection,
  tableNode: TableNode,
): boolean {
  if ($isTableSelection(selection)) {
    const anchorNode = selection.anchor.getNode() as TableCellNode;
    const focusNode = selection.focus.getNode() as TableCellNode;
    if (tableNode && anchorNode && focusNode) {
      const [map] = $computeTableMap(tableNode, anchorNode, focusNode);
      return (
        anchorNode.getKey() === map[0][0].cell.getKey() &&
        focusNode.getKey() === map[map.length - 1].at(-1)!.cell.getKey()
      );
    }
  }
  return false;
}

function selectTableCellNode(tableCell: TableCellNode, fromStart: boolean) {
  if (fromStart) {
    tableCell.selectStart();
  } else {
    tableCell.selectEnd();
  }
}

function $addHighlightToDOM(editor: LexicalEditor, cell: TableDOMCell): void {
  const element = cell.elem;
  const editorThemeClasses = editor._config.theme;
  const node = $getNearestNodeFromDOMNode(element);
  invariant(
    $isTableCellNode(node),
    'Expected to find LexicalNode from Table Cell DOMNode',
  );
  addClassNamesToElement(element, editorThemeClasses.tableCellSelected);
}

function $removeHighlightFromDOM(
  editor: LexicalEditor,
  cell: TableDOMCell,
): void {
  const element = cell.elem;
  const node = $getNearestNodeFromDOMNode(element);
  invariant(
    $isTableCellNode(node),
    'Expected to find LexicalNode from Table Cell DOMNode',
  );
  const editorThemeClasses = editor._config.theme;
  removeClassNamesFromElement(element, editorThemeClasses.tableCellSelected);
}

export function $findCellNode(node: LexicalNode): null | TableCellNode {
  const cellNode = $findMatchingParent(node, $isTableCellNode);
  return $isTableCellNode(cellNode) ? cellNode : null;
}

export function $findTableNode(node: LexicalNode): null | TableNode {
  const tableNode = $findMatchingParent(node, $isTableNode);
  return $isTableNode(tableNode) ? tableNode : null;
}

function $getBlockParentIfFirstNode(node: LexicalNode): ElementNode | null {
  for (
    let prevNode = node, currentNode: LexicalNode | null = node;
    currentNode !== null;
    prevNode = currentNode, currentNode = currentNode.getParent()
  ) {
    if ($isElementNode(currentNode)) {
      if (
        currentNode !== prevNode &&
        currentNode.getFirstChild() !== prevNode
      ) {
        // Not the first child or the initial node
        return null;
      } else if (!currentNode.isInline()) {
        return currentNode;
      }
    }
  }
  return null;
}

function $handleHorizontalArrowKeyRangeSelection(
  editor: LexicalEditor,
  event: KeyboardEvent,
  selection: RangeSelection,
  alter: 'extend' | 'move',
  isBackward: boolean,
  tableNode: TableNode,
  tableObserver: TableObserver,
): boolean {
  const initialFocus = $caretFromPoint(
    selection.focus,
    isBackward ? 'previous' : 'next',
  );
  if ($isExtendableTextPointCaret(initialFocus)) {
    return false;
  }
  let lastCaret = initialFocus;
  // TableCellNode is the only shadow root we are interested in piercing so
  // we find the last internal caret and then check its parent
  for (const nextCaret of $extendCaretToRange(initialFocus).iterNodeCarets(
    'shadowRoot',
  )) {
    if (!($isSiblingCaret(nextCaret) && $isElementNode(nextCaret.origin))) {
      return false;
    }
    lastCaret = nextCaret;
  }
  const lastCaretParent = lastCaret.getParentAtCaret();
  if (!$isTableCellNode(lastCaretParent)) {
    return false;
  }
  const anchorCell = lastCaretParent;
  const focusCaret = $findNextTableCell(
    $getSiblingCaret(anchorCell, lastCaret.direction),
  );
  const anchorCellTable = $findMatchingParent(anchorCell, $isTableNode);
  if (!(anchorCellTable && anchorCellTable.is(tableNode))) {
    return false;
  }
  const anchorCellDOM = editor.getElementByKey(anchorCell.getKey());
  const anchorDOMCell = getDOMCellFromTarget(anchorCellDOM);
  if (!anchorCellDOM || !anchorDOMCell) {
    return false;
  }

  const anchorCellTableElement = $getElementForTableNode(
    editor,
    anchorCellTable,
  );
  tableObserver.table = anchorCellTableElement;
  if (!focusCaret) {
    if (alter === 'extend') {
      // extend the selection from a range inside the cell to a table selection of the cell
      tableObserver.$setAnchorCellForSelection(anchorDOMCell);
      tableObserver.$setFocusCellForSelection(anchorDOMCell, true);
    } else {
      // exit the table
      const outerFocusCaret = $getTableExitCaret(
        $getSiblingCaret(anchorCellTable, initialFocus.direction),
      );
      $setPointFromCaret(selection.anchor, outerFocusCaret);
      $setPointFromCaret(selection.focus, outerFocusCaret);
    }
  } else if (alter === 'extend') {
    const focusDOMCell = getDOMCellFromTarget(
      editor.getElementByKey(focusCaret.origin.getKey()),
    );
    if (!focusDOMCell) {
      return false;
    }
    tableObserver.$setAnchorCellForSelection(anchorDOMCell);
    tableObserver.$setFocusCellForSelection(focusDOMCell, true);
  } else {
    // alter === 'move'
    const innerFocusCaret = $normalizeCaret(focusCaret);
    $setPointFromCaret(selection.anchor, innerFocusCaret);
    $setPointFromCaret(selection.focus, innerFocusCaret);
  }
  stopEvent(event);
  return true;
}

function $getTableExitCaret<D extends CaretDirection>(
  initialCaret: SiblingCaret<TableNode, D>,
): PointCaret<D> {
  const adjacent = $getAdjacentChildCaret(initialCaret);
  return $isChildCaret(adjacent) ? $normalizeCaret(adjacent) : initialCaret;
}

function $findNextTableCell<D extends CaretDirection>(
  initialCaret: SiblingCaret<TableCellNode, D>,
): null | ChildCaret<TableCellNode, D> {
  for (const nextCaret of $extendCaretToRange(initialCaret).iterNodeCarets(
    'root',
  )) {
    const {origin} = nextCaret;
    if ($isTableCellNode(origin)) {
      // not sure why ts isn't narrowing here (even if the guard is on nextCaret.origin)
      // but returning a new caret is fine
      if ($isChildCaret(nextCaret)) {
        return $getChildCaret(origin, initialCaret.direction);
      }
    } else if (!$isTableRowNode(origin)) {
      break;
    }
  }
  return null;
}

function $handleArrowKey(
  editor: LexicalEditor,
  event: KeyboardEvent,
  direction: Direction,
  tableNode: TableNode,
  tableObserver: TableObserver,
  tableObservers: TableObservers,
): boolean {
  if (
    (direction === 'up' || direction === 'down') &&
    isTypeaheadMenuInView(editor)
  ) {
    return false;
  }

  const selection = $getSelection();

  // Handle arrow key into a table (including from a table into a nested table)
  if (!$isSelectionInTable(selection, tableNode)) {
    if ($isRangeSelection(selection)) {
      if (direction === 'backward') {
        if (selection.focus.offset > 0) {
          return false;
        }
        const parentNode = $getBlockParentIfFirstNode(
          selection.focus.getNode(),
        );
        if (!parentNode) {
          return false;
        }
        const siblingNode = parentNode.getPreviousSibling();
        if (!$isTableNode(siblingNode)) {
          return false;
        }
        stopEvent(event);
        if (event.shiftKey) {
          selection.focus.set(
            siblingNode.getParentOrThrow().getKey(),
            siblingNode.getIndexWithinParent(),
            'element',
          );
        } else {
          siblingNode.selectEnd();
        }
        return true;
      } else if (
        event.shiftKey &&
        (direction === 'up' || direction === 'down')
      ) {
        const focusNode = selection.focus.getNode();
        const isTableUnselect =
          !selection.isCollapsed() &&
          ((direction === 'up' && !selection.isBackward()) ||
            (direction === 'down' && selection.isBackward()));
        if (isTableUnselect) {
          let focusParentNode = $findMatchingParent(focusNode, n =>
            $isTableNode(n),
          );
          if ($isTableCellNode(focusParentNode)) {
            focusParentNode = $findMatchingParent(
              focusParentNode,
              $isTableNode,
            );
          }
          if (focusParentNode !== tableNode) {
            return false;
          }
          if (!focusParentNode) {
            return false;
          }
          const sibling =
            direction === 'down'
              ? focusParentNode.getNextSibling()
              : focusParentNode.getPreviousSibling();
          if (!sibling) {
            return false;
          }
          let newOffset = 0;
          if (direction === 'up') {
            if ($isElementNode(sibling)) {
              newOffset = sibling.getChildrenSize();
            }
          }
          let newFocusNode = sibling;
          if (direction === 'up') {
            if ($isElementNode(sibling)) {
              const lastCell = sibling.getLastChild();
              newFocusNode = lastCell ? lastCell : sibling;
              newOffset = $isTextNode(newFocusNode)
                ? newFocusNode.getTextContentSize()
                : 0;
            }
          }
          const newSelection = selection.clone();

          newSelection.focus.set(
            newFocusNode.getKey(),
            newOffset,
            $isTextNode(newFocusNode) ? 'text' : 'element',
          );
          $setSelection(newSelection);
          stopEvent(event);
          return true;
        } else if ($isRootOrShadowRoot(focusNode)) {
          const selectedNode =
            direction === 'up'
              ? selection.getNodes()[selection.getNodes().length - 1]
              : selection.getNodes()[0];
          if (selectedNode) {
            const tableCellNode = $findParentTableCellNodeInTable(
              tableNode,
              selectedNode,
            );
            if (tableCellNode !== null) {
              const firstDescendant = tableNode.getFirstDescendant();
              const lastDescendant = tableNode.getLastDescendant();
              if (!firstDescendant || !lastDescendant) {
                return false;
              }
              const [firstCellNode] = $getNodeTriplet(firstDescendant);
              const [lastCellNode] = $getNodeTriplet(lastDescendant);
              const firstCellCoords = tableNode.getCordsFromCellNode(
                firstCellNode,
                tableObserver.table,
              );
              const lastCellCoords = tableNode.getCordsFromCellNode(
                lastCellNode,
                tableObserver.table,
              );
              const firstCellDOM = tableNode.getDOMCellFromCordsOrThrow(
                firstCellCoords.x,
                firstCellCoords.y,
                tableObserver.table,
              );
              const lastCellDOM = tableNode.getDOMCellFromCordsOrThrow(
                lastCellCoords.x,
                lastCellCoords.y,
                tableObserver.table,
              );
              tableObserver.$setAnchorCellForSelection(firstCellDOM);
              tableObserver.$setFocusCellForSelection(lastCellDOM, true);
              return true;
            }
          }
          return false;
        } else {
          let focusParentNode = $findMatchingParent(
            focusNode,
            n => $isElementNode(n) && !n.isInline(),
          );
          if ($isTableCellNode(focusParentNode)) {
            focusParentNode = $findMatchingParent(
              focusParentNode,
              $isTableNode,
            );
          }
          if (!focusParentNode) {
            return false;
          }
          const sibling =
            direction === 'down'
              ? focusParentNode.getNextSibling()
              : focusParentNode.getPreviousSibling();
          if (
            $isTableNode(sibling) &&
            tableObserver.tableNodeKey === sibling.getKey()
          ) {
            const firstDescendant = sibling.getFirstDescendant();
            const lastDescendant = sibling.getLastDescendant();
            if (!firstDescendant || !lastDescendant) {
              return false;
            }
            const [firstCellNode] = $getNodeTriplet(firstDescendant);
            const [lastCellNode] = $getNodeTriplet(lastDescendant);
            const newSelection = selection.clone();
            newSelection.focus.set(
              (direction === 'up' ? firstCellNode : lastCellNode).getKey(),
              direction === 'up' ? 0 : lastCellNode.getChildrenSize(),
              'element',
            );
            stopEvent(event);
            $setSelection(newSelection);
            return true;
          }
        }
      }
    }
    if (direction === 'down' && $isScrollableTablesActive(editor)) {
      // Enable Firefox workaround
      tableObservers.setShouldCheckSelectionForTable(tableNode.getKey());
    }
    return false;
  }

  if ($isRangeSelection(selection)) {
    if (direction === 'backward' || direction === 'forward') {
      const alter = event.shiftKey ? 'extend' : 'move';
      return $handleHorizontalArrowKeyRangeSelection(
        editor,
        event,
        selection,
        alter,
        direction === 'backward',
        tableNode,
        tableObserver,
      );
    }

    if (selection.isCollapsed()) {
      const {anchor, focus} = selection;
      const anchorCellNode = $findMatchingParent(
        anchor.getNode(),
        $isTableCellNode,
      );
      const focusCellNode = $findMatchingParent(
        focus.getNode(),
        $isTableCellNode,
      );
      if (
        !$isTableCellNode(anchorCellNode) ||
        !anchorCellNode.is(focusCellNode)
      ) {
        return false;
      }
      const anchorCellTable = $findTableNode(anchorCellNode);
      if (anchorCellTable !== tableNode && anchorCellTable != null) {
        const anchorCellTableElement = getTableElement(
          anchorCellTable,
          editor.getElementByKey(anchorCellTable.getKey()),
        );
        if (anchorCellTableElement != null) {
          tableObserver.table = getTable(
            anchorCellTable,
            anchorCellTableElement,
          );
          return $handleArrowKey(
            editor,
            event,
            direction,
            anchorCellTable,
            tableObserver,
            tableObservers,
          );
        }
      }

      const anchorCellDom = editor.getElementByKey(anchorCellNode.__key);
      const anchorDOM = editor.getElementByKey(anchor.key);
      if (anchorDOM == null || anchorCellDom == null) {
        return false;
      }

      let edgeSelectionRect;
      if (anchor.type === 'element') {
        edgeSelectionRect = anchorDOM.getBoundingClientRect();
      } else {
        const domSelection = getDOMSelection(getEditorWindow(editor));
        if (domSelection === null || domSelection.rangeCount === 0) {
          return false;
        }

        const range = domSelection.getRangeAt(0);
        edgeSelectionRect = range.getBoundingClientRect();
      }

      const edgeChild =
        direction === 'up'
          ? anchorCellNode.getFirstChild()
          : anchorCellNode.getLastChild();
      if (edgeChild == null) {
        return false;
      }

      const edgeChildDOM = editor.getElementByKey(edgeChild.__key);

      if (edgeChildDOM == null) {
        return false;
      }

      const edgeRect = edgeChildDOM.getBoundingClientRect();
      const isExiting =
        direction === 'up'
          ? edgeRect.top > edgeSelectionRect.top - edgeSelectionRect.height
          : edgeSelectionRect.bottom + edgeSelectionRect.height >
            edgeRect.bottom;

      if (isExiting) {
        stopEvent(event);

        const cords = tableNode.getCordsFromCellNode(
          anchorCellNode,
          tableObserver.table,
        );

        if (event.shiftKey) {
          const cell = tableNode.getDOMCellFromCordsOrThrow(
            cords.x,
            cords.y,
            tableObserver.table,
          );
          tableObserver.$setAnchorCellForSelection(cell);
          tableObserver.$setFocusCellForSelection(cell, true);
        } else {
          return selectTableNodeInDirection(
            tableObserver,
            tableNode,
            cords.x,
            cords.y,
            direction,
          );
        }

        return true;
      }
    }
  } else if ($isTableSelection(selection)) {
    const {anchor, focus, tableKey} = selection;
    if (tableKey !== tableNode.getKey()) {
      return false;
    }
    const anchorCellNode = $findMatchingParent(
      anchor.getNode(),
      $isTableCellNode,
    );
    const focusCellNode = $findMatchingParent(
      focus.getNode(),
      $isTableCellNode,
    );

    const [tableNodeFromSelection] = selection.getNodes();
    invariant(
      $isTableNode(tableNodeFromSelection),
      '$handleArrowKey: TableSelection.getNodes()[0] expected to be TableNode',
    );
    const tableElement = getTableElement(
      tableNodeFromSelection,
      editor.getElementByKey(tableNodeFromSelection.getKey()),
    );
    if (
      !$isTableCellNode(anchorCellNode) ||
      !$isTableCellNode(focusCellNode) ||
      !$isTableNode(tableNodeFromSelection) ||
      tableElement == null
    ) {
      return false;
    }
    tableObserver.$updateTableTableSelection(selection);

    const grid = getTable(tableNodeFromSelection, tableElement);
    const cordsAnchor = tableNode.getCordsFromCellNode(anchorCellNode, grid);
    const anchorCell = tableNode.getDOMCellFromCordsOrThrow(
      cordsAnchor.x,
      cordsAnchor.y,
      grid,
    );
    tableObserver.$setAnchorCellForSelection(anchorCell);

    stopEvent(event);

    if (event.shiftKey) {
      const [tableMap, anchorValue, focusValue] = $computeTableMap(
        tableNode,
        anchorCellNode,
        focusCellNode,
      );
      return $adjustFocusInDirection(
        tableObserver,
        tableMap,
        anchorValue,
        focusValue,
        direction,
      );
    } else {
      focusCellNode.selectEnd();
    }

    return true;
  }

  return false;
}

function stopEvent(event: Event) {
  event.preventDefault();
  event.stopImmediatePropagation();
  event.stopPropagation();
}

function isTypeaheadMenuInView(editor: LexicalEditor) {
  // There is no inbuilt way to check if the component picker is in view
  // but we can check if the root DOM element has the aria-controls attribute "typeahead-menu".
  const root = editor.getRootElement();
  if (!root) {
    return false;
  }
  return (
    root.hasAttribute('aria-controls') &&
    root.getAttribute('aria-controls') === 'typeahead-menu'
  );
}

function $insertParagraphAtTableEdge(
  edgePosition: 'first' | 'last',
  tableNode: TableNode,
  children?: LexicalNode[],
) {
  const paragraphNode = $createParagraphNode();
  if (edgePosition === 'first') {
    tableNode.insertBefore(paragraphNode);
  } else {
    tableNode.insertAfter(paragraphNode);
  }
  paragraphNode.append(...(children || []));
  paragraphNode.selectEnd();
}

function $getTableEdgeCursorPosition(
  editor: LexicalEditor,
  selection: RangeSelection,
  tableNode: TableNode,
) {
  const tableNodeParent = tableNode.getParent();
  if (!tableNodeParent) {
    return undefined;
  }

  // TODO: Add support for nested tables
  const domSelection = getDOMSelection(getEditorWindow(editor));
  if (!domSelection) {
    return undefined;
  }
  const domAnchorNode = domSelection.anchorNode;
  const tableNodeParentDOM = editor.getElementByKey(tableNodeParent.getKey());
  const tableElement = getTableElement(
    tableNode,
    editor.getElementByKey(tableNode.getKey()),
  );
  // We are only interested in the scenario where the
  // native selection anchor is:
  // - at or inside the table's parent DOM
  // - and NOT at or inside the table DOM
  // It may be adjacent to the table DOM (e.g. in a wrapper)
  if (
    !domAnchorNode ||
    !tableNodeParentDOM ||
    !tableElement ||
    !tableNodeParentDOM.contains(domAnchorNode) ||
    tableElement.contains(domAnchorNode)
  ) {
    return undefined;
  }

  const anchorCellNode = $findMatchingParent(selection.anchor.getNode(), n =>
    $isTableCellNode(n),
  ) as TableCellNode | null;
  if (!anchorCellNode) {
    return undefined;
  }

  const parentTable = $findMatchingParent(anchorCellNode, n => $isTableNode(n));
  if (!$isTableNode(parentTable) || !parentTable.is(tableNode)) {
    return undefined;
  }

  const [tableMap, cellValue] = $computeTableMap(
    tableNode,
    anchorCellNode,
    anchorCellNode,
  );
  const firstCell = tableMap[0][0];
  const lastCell = tableMap[tableMap.length - 1][tableMap[0].length - 1];
  const {startRow, startColumn} = cellValue;

  const isAtFirstCell =
    startRow === firstCell.startRow && startColumn === firstCell.startColumn;
  const isAtLastCell =
    startRow === lastCell.startRow && startColumn === lastCell.startColumn;

  if (isAtFirstCell) {
    return 'first';
  } else if (isAtLastCell) {
    return 'last';
  } else {
    return undefined;
  }
}

export function $getObserverCellFromCellNodeOrThrow(
  tableObserver: TableObserver,
  tableCellNode: TableCellNode,
): TableDOMCell {
  const {tableNode} = tableObserver.$lookup();
  const currentCords = tableNode.getCordsFromCellNode(
    tableCellNode,
    tableObserver.table,
  );
  return tableNode.getDOMCellFromCordsOrThrow(
    currentCords.x,
    currentCords.y,
    tableObserver.table,
  );
}

export function $getNearestTableCellInTableFromDOMNode(
  tableNode: TableNode,
  startingDOM: Node,
  editorState?: EditorState,
) {
  return $findParentTableCellNodeInTable(
    tableNode,
    $getNearestNodeFromDOMNode(startingDOM, editorState),
  );
}