// @flow
import { useRef, useLayoutEffect, useMemo } from 'react';

const getDocument = (): Document | void =>
  typeof document !== 'undefined' ? document : undefined;

const DEFAULT_REF = { current: null };

export const useFocusTrap = ({
  show,
  refToRestoreFocus = DEFAULT_REF,
}: {|
  show: boolean,
  refToRestoreFocus?: { current: HTMLElement | null },
|}) => {
  const modalRef = useRef<HTMLDivElement | null>(null);
  const computedRefToRestoreFocus = useRef(refToRestoreFocus.current);
  const isVisible = useRef(false);
  const childFocusSet = useRef();

  useMemo(() => {
    // gets the documents current active element so we can go back to focusing on it on close
    // we do this on memo because this has to happen before render (& before useEffects). This won't always restore 100%
    // as useMemo isn't guaranteed to always memoize
    const documentElement = getDocument();
    if (show && documentElement && !computedRefToRestoreFocus.current) {
      computedRefToRestoreFocus.current = documentElement.activeElement;
    }
  }, [show]);

  useLayoutEffect(() => {
    // on mount, we listen for a focus event, if there is one we assume a child is trying to focus, so we
    // refocus the child
    const documentElement = getDocument();
    function trapFocus({
      relatedTarget: previousFocusTarget,
      target,
    }: FocusEvent) {
      const { current: modal } = modalRef;

      // for the first focus
      if (
        !childFocusSet.current &&
        modal &&
        // related target is the target we're focusing away from. If it exists inside the modal, we should re-focus it
        previousFocusTarget instanceof HTMLElement &&
        modal.contains(previousFocusTarget) &&
        // parent.contains(parent) is true, so we don't want to double focus
        modal !== previousFocusTarget
      ) {
        childFocusSet.current = true;

        // Refocus to the previous focus target if and only if the new target is *not* another
        // child of the modal.
        if (
          previousFocusTarget &&
          target &&
          // $FlowFixMe target isn't a Node as required by contains.
          (!modal.contains(target) || modal === target)
        ) {
          previousFocusTarget.focus();
        }
      }

      // for all focuses after - we trap the focus so they can't shift-tab outside the modal
      if (
        childFocusSet.current &&
        target &&
        modal &&
        target instanceof HTMLElement &&
        !modal.contains(target)
      ) {
        modal.focus();
      }
    }

    if (show && documentElement && !childFocusSet.current) {
      documentElement.addEventListener('focus', trapFocus, true);
    }

    return () => {
      if (documentElement) {
        documentElement.removeEventListener('focus', trapFocus, true);
      }
    };
  }, [show, modalRef, childFocusSet]);

  useLayoutEffect(() => {
    const isNowVisible = show && !isVisible.current;
    const isNowHidden = !show && isVisible.current;

    if (isNowVisible) {
      isVisible.current = true;

      const { current: modal } = modalRef;
      if (modal) {
        modal.setAttribute('tabIndex', '-1');
        modal.focus();
      }
    }

    if (isNowHidden) {
      isVisible.current = false;

      // try and focus the passed in refToRestoreFocus
      if (refToRestoreFocus.current instanceof HTMLElement) {
        return refToRestoreFocus.current.focus();
      }

      // otherwise, use the computed one
      if (computedRefToRestoreFocus.current instanceof HTMLElement) {
        return computedRefToRestoreFocus.current.focus();
      }
    }
  }, [show, modalRef, refToRestoreFocus]);

  return modalRef;
};
