import { useRef } from "react";

// Overlays opened from state have no Radix trigger to refocus, so remember what had focus when one
// opened and return there on close, unless the caller handled it.
export const useReturnFocus = (
  onOpenAutoFocus?: (event: Event) => void,
  onCloseAutoFocus?: (event: Event) => void,
) => {
  const returnTo = useRef<HTMLElement | null>(null);
  return {
    onOpenAutoFocus: (event: Event) => {
      returnTo.current = document.activeElement instanceof HTMLElement ? document.activeElement : null;
      onOpenAutoFocus?.(event);
    },
    onCloseAutoFocus: (event: Event) => {
      onCloseAutoFocus?.(event);
      const target = returnTo.current;
      if (event.defaultPrevented || !target || !target.isConnected || target === document.body) return;
      event.preventDefault();
      target.focus();
    },
  };
};