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

interface ItemTops {
  [id: string]: number | undefined;
}

interface ListItemRefsById {
  [id: string]: HTMLLIElement | undefined;
}

interface ItemOffsets {
  [id: string]: number | undefined;
}

interface ItemStyles {
  position: 'relative';
  top: number;
  transition: string;
}

interface UseAnimatedListItemArgs {
  /**
   * **IMPORTANT - This list must be memoized. The array reference should only change when the
   * list of items changes.**
   *
   * The list of item keys. Each key should uniquely identify a `<ListItem>` component, and
   * should be the same value passed to the `ListItem`'s `key` prop.
   */
  keys: string[] | undefined;
}

interface UseAnimatedListItemReturns {
  /**
   * Item styles that should be passed to the `sx` prop of each `ListItem`. The styles are
   * mapped per item key.
   */
  itemStyles: { [key: string]: ItemStyles | undefined };

  /**
   * Function that must be called within the `ref` callback prop of each `ListItem`.
   *
   * @param {string} key  - the `ListItem`'s unique key.
   * @param {HTMLLIElement} li  - the list item element reference.
   * @returns {void}
   */
  updateItemRef: (key: string, li: HTMLLIElement | null) => void;

  /**
   * Function that must be called within each and every callback that causes the list of items
   * to refresh or update. This is used to take a "snapshot" of the items current positions prior
   * to any change.
   *
   * @returns {void}
   */
  updateItemPositions: () => void;
}

/**
 * This hook can be plugged into a MUI <List> to animate the <ListItem> components whenever
 * change their order.
 */
export const useAnimatedListItems = ({
  keys,
}: UseAnimatedListItemArgs): UseAnimatedListItemReturns => {
  const itemRefs = useRef<ListItemRefsById>({});
  const itemTops = useRef<ItemTops>({});
  const [itemOffsets, setItemOffsets] = useState<ItemOffsets>({});

  useLayoutEffect(() => {
    if (!keys) return;
    const newItemOffsets: ItemOffsets = {};
    keys.forEach((key) => {
      const itemRef = itemRefs.current[key];
      if (itemRef) {
        const currentTop = itemRef.getBoundingClientRect().top;
        const prevTop = itemTops.current[key] || 0;
        const offset = -(currentTop - prevTop);
        newItemOffsets[key] = offset;
      }
    });
    setItemOffsets(newItemOffsets);

    requestAnimationFrame(() => {
      setItemOffsets({});
    });
  }, [keys]);

  let itemStyles: UseAnimatedListItemReturns['itemStyles'] = {};
  keys?.forEach((key) => {
    itemStyles[key] = {
      position: 'relative',
      top: itemOffsets[key] || 0,
      transition:
        !itemTops.current[key] || itemOffsets[key] ? 'top 0s' : 'top 1s',
    };
  });

  return {
    itemStyles,
    updateItemRef: (key, li) => {
      li === null ? delete itemRefs.current[key] : (itemRefs.current[key] = li);
    },
    updateItemPositions: () => {
      keys?.forEach((key) => {
        const itemRef = itemRefs.current[key];
        if (itemRef) {
          itemTops.current[key] = itemRef.getBoundingClientRect().top;
        }
      });
    },
  };
};