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

type RowKey = string;
type DocumentId = string | number;

interface SelectionEventLike {
  shiftKey?: boolean;
  metaKey?: boolean;
  ctrlKey?: boolean;
  preventDefault?: () => void;
}

interface UseDocumentSelectionOptions {
  resolveDocumentRowKey: (id: DocumentId | null | undefined) => RowKey | null | undefined;
  resolveFolderRowKey: (id: DocumentId | null | undefined) => RowKey | null | undefined;
  isDocumentRowKey: (key?: RowKey | null) => boolean;
  isFolderRowKey: (key?: RowKey | null) => boolean;
  getRowId: (key?: RowKey | null) => DocumentId | null | undefined;
  initialEntries?: RowKey[];
}

interface ApplySelectionOptions {
  anchor?: RowKey | null;
  interactedKeys?: RowKey[];
}

const DEFAULT_INITIAL_ENTRIES: RowKey[] = [];

export const useDocumentSelection = ({
  resolveDocumentRowKey,
  resolveFolderRowKey,
  isDocumentRowKey,
  isFolderRowKey,
  getRowId,
  initialEntries = DEFAULT_INITIAL_ENTRIES,
}: UseDocumentSelectionOptions) => {
  const [selectedEntries, setSelectedEntries] = useState<RowKey[]>(initialEntries);
  const [selectionOrder, setSelectionOrder] = useState<RowKey[]>(initialEntries);
  const selectionOrderRef = useRef<RowKey[]>(initialEntries);
  const selectionAnchorRef = useRef<RowKey | null>(null);
  const selectionInitializedRef = useRef(false);
  const [focusedDocumentId, setFocusedDocumentId] = useState<DocumentId | null>(null);
  const [focusedRowKey, setFocusedRowKey] = useState<RowKey | null>(null);

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

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

  const updateSelectionOrder = useCallback((nextSelection: RowKey[], interactedKeys: RowKey[] = []) => {
    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: Array<RowKey | null | undefined>,
      { anchor, interactedKeys = [] }: ApplySelectionOptions = {},
    ) => {
      const visibleRowKeySet = visibleRowKeySetRef.current;
      const unique: RowKey[] = [];

      (rowKeys || []).forEach((key) => {
        if (!key) return;
        let canonicalKey: RowKey | null | undefined = 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;
      }

      setSelectedEntries(unique);
      updateSelectionOrder(unique, interactedKeys);

      const nextFocusedDocumentId: DocumentId | null = (() => {
        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((key) => isDocumentRowKey(key)) ?? null;
        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 handleEntrySelection = useCallback(
    (rowKey: RowKey | null | undefined, event?: SelectionEventLike) => {
      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 && selectedEntries.length) {
        anchorKey = selectedEntries[selectedEntries.length - 1];
      }
      if (!anchorKey) {
        anchorKey = rowKey;
      }

      let nextKeys: RowKey[] = [];
      let interactedKeys: RowKey[] = [];

      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(selectedEntries);
          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 (selectedEntries.includes(rowKey)) {
          nextKeys = selectedEntries.filter((key) => key !== rowKey);
          interactedKeys = [];
        } else {
          nextKeys = [...selectedEntries, rowKey];
          interactedKeys = [rowKey];
        }
        anchorKey = rowKey;
      } else {
        nextKeys = [rowKey];
        interactedKeys = [rowKey];
        anchorKey = rowKey;
      }

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

  const promoteSelectionOrder = useCallback(
    (docId?: DocumentId | null) => {
      if (!docId) return;
      const rowKey = resolveDocumentRowKey(docId);
      if (!rowKey) return;

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

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

  return {
    selectedEntries,
    setSelectedEntries,
    selectionOrder,
    setSelectionOrder,
    selectionOrderRef,
    selectionAnchorRef,
    selectionInitializedRef,
    focusedDocumentId,
    setFocusedDocumentId,
    focusedRowKey,
    setFocusedRowKey,
    applySelection,
    clearSelection,
    handleEntrySelection,
    promoteSelectionOrder,
    configureSelectionEnvironment,
  };
};

export default useDocumentSelection;
