import { useCallback, useLayoutEffect, useRef, type MutableRefObject, type UIEvent } from 'react';

interface UseRetainedScrollTopOptions {
  restoreKey: string | null;
  scrollTopRef?: MutableRefObject<number>;
  suspendSave?: boolean;
}

export function useRetainedScrollTop<T extends HTMLElement>({
  restoreKey,
  scrollTopRef,
  suspendSave = false,
}: UseRetainedScrollTopOptions) {
  const fallbackScrollTopRef = useRef(0);
  const savedScrollTopRef = scrollTopRef ?? fallbackScrollTopRef;
  const nodeRef = useRef<T | null>(null);
  const previousRestoreKeyRef = useRef(restoreKey);
  const pendingRestoreTopRef = useRef<number | null>(null);
  const suspendSaveRef = useRef(suspendSave);

  suspendSaveRef.current = suspendSave;

  const scrollRef = useCallback(
    (node: T | null) => {
      const previousNode = nodeRef.current;
      if (!node && previousNode && !suspendSaveRef.current) {
        savedScrollTopRef.current = previousNode.scrollTop;
      }

      nodeRef.current = node;
      if (node) {
        node.scrollTop = pendingRestoreTopRef.current ?? savedScrollTopRef.current;
      }
    },
    [savedScrollTopRef]
  );

  const onScroll = useCallback(
    (event: UIEvent<T>) => {
      if (suspendSaveRef.current) return;
      savedScrollTopRef.current = event.currentTarget.scrollTop;
    },
    [savedScrollTopRef]
  );

  useLayoutEffect(() => {
    if (previousRestoreKeyRef.current !== restoreKey) {
      previousRestoreKeyRef.current = restoreKey;
      pendingRestoreTopRef.current = savedScrollTopRef.current;
    }

    const node = nodeRef.current;
    const pendingRestoreTop = pendingRestoreTopRef.current;
    if (!node || pendingRestoreTop == null) return;

    node.scrollTop = pendingRestoreTop;
    if (!suspendSave) {
      pendingRestoreTopRef.current = null;
    }
  }, [restoreKey, savedScrollTopRef, suspendSave]);

  return { scrollRef, onScroll };
}
