/**
 * 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 { ReactPortal, useCallback, useEffect, useRef, useState } from 'react';
import { $getRoot, $getSelection, $isRangeSelection, DEPRECATED_$isGridSelection } from 'lexical';
import { createPortal } from 'react-dom';
import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext';
import useLexicalEditable from '@lexical/react/useLexicalEditable';
import {
  $deleteTableColumn,
  $getElementGridForTableNode,
  $getTableCellNodeFromLexicalNode,
  $getTableColumnIndexFromTableCellNode,
  $getTableNodeFromLexicalNodeOrThrow,
  $getTableRowIndexFromTableCellNode,
  $insertTableColumn,
  $insertTableRow,
  $isTableCellNode,
  $isTableNode,
  $isTableRowNode,
  $removeTableRowAtIndex,
  getTableSelectionFromTableElement,
  HTMLTableElementWithWithTableSelectionState,
  TableCellHeaderStates,
  TableCellNode,
} from '@lexical/table';

import { DropdownMenu } from '../../../dropdown-menu';
import { IconButton } from '../../../icon-button';
import { Icon } from '../../../iconography';
import { useTr } from '../../i18n';

type TableStats = {
  rows: number;
  columns: number;
};

type TableCellActionMenuProps = Readonly<{
  tableCellNode: TableCellNode;
  tableStats: TableStats;
}>;

function TableActionMenu({ tableCellNode: _tableCellNode, tableStats }: TableCellActionMenuProps) {
  const [editor] = useLexicalComposerContext();
  const [tableCellNode, updateTableCellNode] = useState(_tableCellNode);
  const [selectionCounts, updateSelectionCounts] = useState({
    columns: 1,
    rows: 1,
  });
  const tr = useTr();

  useEffect(() => {
    return editor.registerMutationListener(TableCellNode, nodeMutations => {
      const nodeUpdated = nodeMutations.get(tableCellNode.getKey()) === 'updated';

      if (nodeUpdated) {
        editor.getEditorState().read(() => {
          updateTableCellNode(tableCellNode.getLatest());
        });
      }
    });
  }, [editor, tableCellNode]);

  useEffect(() => {
    editor.getEditorState().read(() => {
      const selection = $getSelection();

      if (DEPRECATED_$isGridSelection(selection)) {
        const selectionShape = selection.getShape();

        updateSelectionCounts({
          columns: selectionShape.toX - selectionShape.fromX + 1,
          rows: selectionShape.toY - selectionShape.fromY + 1,
        });
      }
    });
  }, [editor]);

  const clearTableSelection = useCallback(() => {
    editor.update(() => {
      if (tableCellNode.isAttached()) {
        const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);
        const tableElement = editor.getElementByKey(tableNode.getKey()) as HTMLTableElementWithWithTableSelectionState;

        if (!tableElement) {
          throw new Error('Expected to find tableElement in DOM');
        }

        const tableSelection = getTableSelectionFromTableElement(tableElement);
        if (tableSelection !== null) {
          tableSelection.clearHighlight();
        }

        tableNode.markDirty();
        updateTableCellNode(tableCellNode.getLatest());
      }

      const rootNode = $getRoot();
      rootNode.selectStart();
    });
  }, [editor, tableCellNode]);

  const insertTableRowAtSelection = useCallback(
    (shouldInsertAfter: boolean) => {
      editor.update(() => {
        const selection = $getSelection();

        const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);

        let tableRowIndex;

        if (DEPRECATED_$isGridSelection(selection)) {
          const selectionShape = selection.getShape();
          tableRowIndex = shouldInsertAfter ? selectionShape.toY : selectionShape.fromY;
        } else {
          tableRowIndex = $getTableRowIndexFromTableCellNode(tableCellNode);
        }

        const grid = $getElementGridForTableNode(editor, tableNode);

        $insertTableRow(tableNode, tableRowIndex, shouldInsertAfter, selectionCounts.rows, grid);

        clearTableSelection();
      });
    },
    [editor, tableCellNode, selectionCounts.rows, clearTableSelection],
  );

  const insertTableColumnAtSelection = useCallback(
    (shouldInsertAfter: boolean) => {
      editor.update(() => {
        const selection = $getSelection();

        const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);

        let tableColumnIndex;

        if (DEPRECATED_$isGridSelection(selection)) {
          const selectionShape = selection.getShape();
          tableColumnIndex = shouldInsertAfter ? selectionShape.toX : selectionShape.fromX;
        } else {
          tableColumnIndex = $getTableColumnIndexFromTableCellNode(tableCellNode);
        }

        const grid = $getElementGridForTableNode(editor, tableNode);

        $insertTableColumn(tableNode, tableColumnIndex, shouldInsertAfter, selectionCounts.columns, grid);

        clearTableSelection();
      });
    },
    [editor, tableCellNode, selectionCounts.columns, clearTableSelection],
  );

  const deleteTableRowAtSelection = useCallback(() => {
    editor.update(() => {
      const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);
      const tableRowIndex = $getTableRowIndexFromTableCellNode(tableCellNode);

      $removeTableRowAtIndex(tableNode, tableRowIndex);

      clearTableSelection();
    });
  }, [editor, tableCellNode, clearTableSelection]);

  const deleteTableAtSelection = useCallback(() => {
    editor.update(() => {
      const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);
      tableNode.remove();

      clearTableSelection();
    });
  }, [editor, tableCellNode, clearTableSelection]);

  const deleteTableColumnAtSelection = useCallback(() => {
    editor.update(() => {
      const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);

      const tableColumnIndex = $getTableColumnIndexFromTableCellNode(tableCellNode);

      $deleteTableColumn(tableNode, tableColumnIndex);

      clearTableSelection();
    });
  }, [editor, tableCellNode, clearTableSelection]);

  const toggleTableRowIsHeader = useCallback(() => {
    editor.update(() => {
      const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);

      const tableRowIndex = $getTableRowIndexFromTableCellNode(tableCellNode);

      const tableRows = tableNode.getChildren();

      if (tableRowIndex >= tableRows.length || tableRowIndex < 0) {
        throw new Error('Expected table cell to be inside of table row.');
      }

      const tableRow = tableRows[tableRowIndex];

      if (!$isTableRowNode(tableRow)) {
        throw new Error('Expected table row');
      }

      tableRow.getChildren().forEach(tableCell => {
        if (!$isTableCellNode(tableCell)) {
          throw new Error('Expected table cell');
        }

        tableCell.toggleHeaderStyle(TableCellHeaderStates.ROW);
      });

      clearTableSelection();
    });
  }, [editor, tableCellNode, clearTableSelection]);

  const toggleTableColumnIsHeader = useCallback(() => {
    editor.update(() => {
      const tableNode = $getTableNodeFromLexicalNodeOrThrow(tableCellNode);

      const tableColumnIndex = $getTableColumnIndexFromTableCellNode(tableCellNode);

      const tableRows = tableNode.getChildren();

      for (let r = 0; r < tableRows.length; r++) {
        const tableRow = tableRows[r];

        if (!$isTableRowNode(tableRow)) {
          throw new Error('Expected table row');
        }

        const tableCells = tableRow.getChildren();

        if (tableColumnIndex >= tableCells.length || tableColumnIndex < 0) {
          throw new Error('Expected table cell to be inside of table row.');
        }

        const tableCell = tableCells[tableColumnIndex];

        if (!$isTableCellNode(tableCell)) {
          throw new Error('Expected table cell');
        }

        tableCell.toggleHeaderStyle(TableCellHeaderStates.COLUMN);
      }

      clearTableSelection();
    });
  }, [editor, tableCellNode, clearTableSelection]);

  return (
    <>
      <DropdownMenu.Item onSelect={() => insertTableRowAtSelection(false)}>
        {tr('actionTableInsertRowsAbove', selectionCounts.rows)}
      </DropdownMenu.Item>
      <DropdownMenu.Item onSelect={() => insertTableRowAtSelection(true)}>
        {tr('actionTableInsertRowsBelow', selectionCounts.rows)}
      </DropdownMenu.Item>
      <DropdownMenu.Item onSelect={() => insertTableColumnAtSelection(false)}>
        {tr('actionTableInsertColumnsBefore', selectionCounts.columns)}
      </DropdownMenu.Item>
      <DropdownMenu.Item onSelect={() => insertTableColumnAtSelection(true)}>
        {tr('actionTableInsertColumnsAfter', selectionCounts.columns)}
      </DropdownMenu.Item>
      <DropdownMenu.Item onSelect={() => toggleTableRowIsHeader()}>
        {tr(
          (tableCellNode.__headerState & TableCellHeaderStates.ROW) === TableCellHeaderStates.ROW
            ? 'actionTableRemoveRowHeader'
            : 'actionTableAddRowHeader',
        )}
      </DropdownMenu.Item>
      <DropdownMenu.Item onSelect={() => toggleTableColumnIsHeader()}>
        {tr(
          (tableCellNode.__headerState & TableCellHeaderStates.COLUMN) === TableCellHeaderStates.COLUMN
            ? 'actionTableRemoveColumnHeader'
            : 'actionTableAddColumnHeader',
        )}
      </DropdownMenu.Item>
      <DropdownMenu.Separator />
      {tableStats.columns > 1 && (
        <DropdownMenu.Item onSelect={() => deleteTableColumnAtSelection()}>
          {tr('actionTableDeleteColumn')}
        </DropdownMenu.Item>
      )}
      {tableStats.rows > 1 && (
        <DropdownMenu.Item onSelect={() => deleteTableRowAtSelection()}>{tr('actionTableDeleteRow')}</DropdownMenu.Item>
      )}
      <DropdownMenu.Item onSelect={() => deleteTableAtSelection()}>{tr('actionTableDeleteTable')}</DropdownMenu.Item>
    </>
  );
}

function TableCellActionMenuContainer({ anchorElem }: { anchorElem: HTMLElement }) {
  const [editor] = useLexicalComposerContext();
  const tr = useTr();

  const menuButtonRef = useRef(null);

  const [isMenuOpen, setIsMenuOpen] = useState(false);
  const [tableCellNode, setTableMenuCellNode] = useState<TableCellNode | null>(null);
  const [tableStats, setTablestats] = useState<TableStats>({ rows: 1, columns: 1 });

  const moveMenu = useCallback(() => {
    if (isMenuOpen) {
      return;
    }

    const menu = menuButtonRef.current;
    const selection = $getSelection();
    const nativeSelection = window.getSelection();
    const activeElement = document.activeElement;

    if (selection == null || menu == null) {
      setTableMenuCellNode(null);
      return;
    }

    const rootElement = editor.getRootElement();

    if (
      $isRangeSelection(selection) &&
      rootElement !== null &&
      nativeSelection !== null &&
      rootElement.contains(nativeSelection.anchorNode)
    ) {
      const tableCellNodeFromSelection = $getTableCellNodeFromLexicalNode(selection.anchor.getNode());

      if (tableCellNodeFromSelection == null) {
        setTableMenuCellNode(null);
        return;
      }

      const tableCellParentNodeDOM = editor.getElementByKey(tableCellNodeFromSelection.getKey());

      if (tableCellParentNodeDOM == null) {
        setTableMenuCellNode(null);
        return;
      }

      setTableMenuCellNode(tableCellNodeFromSelection);

      let rows = 1;
      let columns = 1;
      const parentRowNode = tableCellNodeFromSelection.getParent();
      if ($isTableRowNode(parentRowNode)) {
        const parentTableNode = parentRowNode.getParent();
        if ($isTableNode(parentTableNode)) {
          rows = parentTableNode.getChildrenSize();
          columns = parentRowNode.getChildrenSize();
        }
      }

      setTablestats({
        rows,
        columns,
      });
    } else if (!activeElement) {
      setTableMenuCellNode(null);
    }
  }, [editor, isMenuOpen]);

  useEffect(() => {
    return editor.registerUpdateListener(() => {
      editor.getEditorState().read(() => {
        moveMenu();
      });
    });
  });

  useEffect(() => {
    const menuButtonDOM = menuButtonRef.current as HTMLButtonElement | null;

    if (menuButtonDOM != null && tableCellNode != null) {
      const tableCellNodeDOM = editor.getElementByKey(tableCellNode.getKey());

      if (tableCellNodeDOM != null) {
        const tableCellRect = tableCellNodeDOM.getBoundingClientRect();
        const menuRect = menuButtonDOM.getBoundingClientRect();
        const anchorRect = anchorElem.getBoundingClientRect();

        const top = tableCellRect.top - anchorRect.top + 4;
        const left = tableCellRect.right - menuRect.width - 10 - anchorRect.left;

        menuButtonDOM.style.opacity = '1';
        menuButtonDOM.style.transform = `translate(${left}px, ${top}px)`;
      } else {
        menuButtonDOM.style.opacity = '0';
        menuButtonDOM.style.transform = 'translate(-10000px, -10000px)';
      }
    }
  }, [menuButtonRef, tableCellNode, editor, anchorElem]);

  return (
    <div className="c-rte-table-cell-action-button-container" ref={menuButtonRef}>
      {tableCellNode != null && (
        <DropdownMenu.Root
          onOpenChange={isOpen => setIsMenuOpen(isOpen)}
          content={<TableActionMenu tableCellNode={tableCellNode} tableStats={tableStats} />}
        >
          <IconButton size="xs" aria-label={tr('actionTableOpenOptions')}>
            <Icon.Arrow />
          </IconButton>
        </DropdownMenu.Root>
      )}
    </div>
  );
}

export default function TableActionMenuPlugin({
  anchorElem = document.body,
}: {
  anchorElem?: HTMLElement;
}): null | ReactPortal {
  const isEditable = useLexicalEditable();
  return createPortal(isEditable ? <TableCellActionMenuContainer anchorElem={anchorElem} /> : null, anchorElem);
}
