papercrate/frontend/src/app/useDocumentSelection.js

import { useCallback, useRef, useState } from 'react';

const DEFAULT_INITIAL_SELECTION = [];

export const useDocumentSelection = ({
  resolveDocumentRowKey,
  resolveFolderRowKey,
  isDocumentRowKey,
  isFolderRowKey,
  getRowId,
  initialSelection = DEFAULT_INITIAL_SELECTION,
}) => {
  const [selectedRowKeys, setSelectedRowKeys] = useState(initialSelection);
  const [selectionOrder, setSelectionOrder] = useState(initialSelection);
  const selectionOrderRef = useRef(initialSelection);
  const selectionAnchorRef = useRef(null);
  const selectionInitializedRef = useRef(false);
  const [focusedDocumentId, setFocusedDocumentId] = useState(null);
  const [focusedRowKey, setFocusedRowKey] = useState(null);

  const visibleRowKeySetRef = useRef(new Set());
  const navigableRowKeysRef = useRef([]);

  const configureSelectionEnvironment = useCallback(({
    visibleRowKeySet,
    navigableRowKeys,
  }) => {
    if (visibleRowKeySet) {
      visibleRowKeySetRef.current = visibleRowKeySet;
    }
    if (Array.isArray(navigableRowKeys)) {
      navigableRowKeysRef.current = navigableRowKeys;
    }
  }, []);

  const updateSelectionOrder = useCallback((nextSelection, interactedKeys = []) => {
    const nextSet = new Set(nextSelection);
    const previousOrder = selectionOrderRef.current.filter((id) => nextSet.has(id));
    const interacted = (interactedKeys || []).filter((id, index, array) => array.indexOf(id) === index);

    const base = previousOrder.filter((id) => !interacted.includes(id));
    const result = [...base];

    interacted.forEach((id) => {
      if (nextSet.has(id) && !result.includes(id)) {
        result.push(id);
      }
    });

    nextSelection.forEach((id) => {
      if (!result.includes(id)) {
        result.push(id);
      }
    });

    if (
      result.length !== selectionOrderRef.current.length
      || result.some((id, index) => selectionOrderRef.current[index] !== id)
    ) {
      selectionOrderRef.current = result;
      setSelectionOrder(result);
    } else {
      selectionOrderRef.current = result;
    }
  }, []);

  const applySelection = useCallback(
    (rowKeys, { anchor, interactedKeys = [] } = {}) => {
      const visibleRowKeySet = visibleRowKeySetRef.current;
      const unique = [];

      (rowKeys || []).forEach((key) => {
        if (!key) return;
        let canonicalKey = null;
        if (visibleRowKeySet.has(key)) {
          canonicalKey = key;
        } else if (isDocumentRowKey(key)) {
          const id = getRowId(key);
          canonicalKey = id ? resolveDocumentRowKey(id) : null;
        } else if (isFolderRowKey(key)) {
          const id = getRowId(key);
          canonicalKey = id ? resolveFolderRowKey(id) : null;
        }

        if (!canonicalKey || !visibleRowKeySet.has(canonicalKey)) {
          return;
        }

        if (!unique.includes(canonicalKey)) {
          unique.push(canonicalKey);
        }
      });

      let resolvedAnchor = anchor;
      if (resolvedAnchor && !unique.includes(resolvedAnchor)) {
        resolvedAnchor = null;
      }

      setSelectedRowKeys(unique);
      updateSelectionOrder(unique, interactedKeys);

      const nextFocusedDocumentId = (() => {
        if (focusedDocumentId) {
          const focusKey = resolveDocumentRowKey(focusedDocumentId);
          if (focusKey && unique.includes(focusKey)) {
            return focusedDocumentId;
          }
        }

        if (resolvedAnchor && isDocumentRowKey(resolvedAnchor)) {
          return getRowId(resolvedAnchor) || null;
        }

        const lastDocKey = [...unique].reverse().find(isDocumentRowKey);
        return lastDocKey ? getRowId(lastDocKey) || null : null;
      })();

      setFocusedDocumentId(nextFocusedDocumentId);

      if (resolvedAnchor) {
        selectionAnchorRef.current = resolvedAnchor;
      } else if (!unique.length) {
        selectionAnchorRef.current = null;
      } else if (!selectionAnchorRef.current || !unique.includes(selectionAnchorRef.current)) {
        selectionAnchorRef.current = unique[unique.length - 1];
      }

      return { selection: unique, focusKey: selectionAnchorRef.current };
    },
    [
      focusedDocumentId,
      getRowId,
      isDocumentRowKey,
      isFolderRowKey,
      resolveDocumentRowKey,
      resolveFolderRowKey,
      updateSelectionOrder,
    ],
  );

  const clearSelection = useCallback(() => {
    setFocusedRowKey(null);
    applySelection([], { anchor: null, interactedKeys: [] });
  }, [applySelection]);

  const handleRowSelection = useCallback(
    (rowKey, event) => {
      const visibleRowKeySet = visibleRowKeySetRef.current;
      const navigableRowKeys = navigableRowKeysRef.current;
      if (!rowKey || !visibleRowKeySet.has(rowKey)) {
        return;
      }

      setFocusedRowKey(rowKey);

      const shiftKey = Boolean(event?.shiftKey);
      const metaKey = Boolean(event?.metaKey);
      const ctrlKey = Boolean(event?.ctrlKey);
      const additive = metaKey || ctrlKey;

      if (shiftKey) {
        event?.preventDefault?.();
      }

      let anchorKey = selectionAnchorRef.current;
      if (!anchorKey && shiftKey && selectedRowKeys.length) {
        anchorKey = selectedRowKeys[selectedRowKeys.length - 1];
      }
      if (!anchorKey) {
        anchorKey = rowKey;
      }

      let nextKeys = [];
      let interactedKeys = [];

      if (shiftKey && anchorKey) {
        const anchorIndex = navigableRowKeys.indexOf(anchorKey);
        const targetIndex = navigableRowKeys.indexOf(rowKey);
        if (anchorIndex !== -1 && targetIndex !== -1) {
          const [start, end] = anchorIndex <= targetIndex
            ? [anchorIndex, targetIndex]
            : [targetIndex, anchorIndex];
          const range = navigableRowKeys.slice(start, end + 1);
          nextKeys = range;

          const previousSet = new Set(selectedRowKeys);
          interactedKeys = range.filter((key) => key === rowKey || !previousSet.has(key));
          if (!interactedKeys.includes(rowKey)) {
            interactedKeys.push(rowKey);
          }
        } else {
          nextKeys = [rowKey];
          interactedKeys = [rowKey];
        }
      } else if (additive) {
        if (selectedRowKeys.includes(rowKey)) {
          nextKeys = selectedRowKeys.filter((key) => key !== rowKey);
          interactedKeys = [];
        } else {
          nextKeys = [...selectedRowKeys, rowKey];
          interactedKeys = [rowKey];
        }
        anchorKey = rowKey;
      } else {
        nextKeys = [rowKey];
        interactedKeys = [rowKey];
        anchorKey = rowKey;
      }

      applySelection(nextKeys, { anchor: anchorKey, interactedKeys });
    },
    [applySelection, selectedRowKeys],
  );

  const promoteSelectionOrder = useCallback(
    (docId) => {
      if (!docId) return;
      const rowKey = resolveDocumentRowKey(docId);
      if (!rowKey) return;

      if (!selectedRowKeys.includes(rowKey)) {
        return;
      }

      updateSelectionOrder(selectedRowKeys, [rowKey]);
    },
    [resolveDocumentRowKey, selectedRowKeys, updateSelectionOrder],
  );

  return {
    selectedRowKeys,
    setSelectedRowKeys,
    selectionOrder,
    setSelectionOrder,
    selectionOrderRef,
    selectionAnchorRef,
    selectionInitializedRef,
    focusedDocumentId,
    setFocusedDocumentId,
    focusedRowKey,
    setFocusedRowKey,
    applySelection,
    clearSelection,
    handleRowSelection,
    promoteSelectionOrder,
    configureSelectionEnvironment,
  };
};

export default useDocumentSelection;