import { FocusScope , Keys , createFocusableTreeWalker , isNil , useEventCallback , useRefState } from "../../shared" ; import { KeyboardEvent , useLayoutEffect } from "react" ; import { isElementInViewport } from "./isElementInViewport" ; export interface UseRestoreFocusOptions { isDisabled ? : boolean ; } function isTabbable ( element ) { return element . getAttribute ( "tabindex" ) !== "-1" ; } export function useRestoreFocus ( focusScope , { isDisabled } = { } ) { const [ elementToRestoreRef , setElementToRestore ] = useRefState < HTMLElement > ( ) ; useLayoutEffect ( ( ) => { setElementToRestore ( document . activeElement as HTMLElement ) ; } , [ isDisabled , setElementToRestore ] ) ; const handleKeyDown = useEventCallback ( ( event ) => { if ( event . key === Keys . tab ) { if ( ! event . isPropagationStopped ( ) ) { const currentActiveElement = event . target ; if ( focusScope . isInScope ( currentActiveElement as HTMLElement ) ) { const walker = createFocusableTreeWalker ( document . body , { acceptElement : isTabbable } ) ; walker . currentNode = currentActiveElement as Node ; const next = ( ) => { return event . shiftKey ? walker . previousNode ( ) : walker . nextNode ( ) as HTMLElement ; } ; let nextElement = next ( ) as HTMLElement ; if ( ! focusScope . isInScope ( nextElement ) ) { const elementToRestore = elementToRestoreRef . current ; if ( document . body . contains ( elementToRestore ) ) { walker . currentNode = elementToRestore ; do { nextElement = next ( ) as HTMLElement ; } while ( nextElement === event . currentTarget || focusScope . isInScope ( nextElement ) ) ; event . preventDefault ( ) ; if ( ! isNil ( nextElement ) ) { nextElement . focus ( ) ; } else { elementToRestore . focus ( ) ; } } } } } } } ) ; useLayoutEffect ( ( ) => { if ( ! isDisabled ) { return ( ) => { if ( focusScope . isInScope ( document . activeElement as HTMLElement ) ) { const elementToRestore = elementToRestoreRef . current ; requestAnimationFrame ( ( ) => { if ( document . body . contains ( elementToRestore ) ) { if ( isElementInViewport ( elementToRestore ) ) { elementToRestore . focus ( ) ; } } } ) ; } } ; } } , [ focusScope , isDisabled , elementToRestoreRef ] ) ; return isDisabled ? { } : { onKeyDown : handleKeyDown } ; }