import React, {
  forwardRef,
  useCallback,
  useEffect,
  useRef,
  useState,
} from "react";
import { useClientLayoutEffect } from "../../../utils-external";
import {
  hideNonTargetElements,
  ownerDocument,
  resolveRef,
} from "../../helpers";
import { focusElement, getTabbableCandidates } from "../../helpers/focus";
import { useMergeRefs } from "../../hooks";
import { useValueAsRef } from "../../hooks/useValueAsRef";
import { Slot } from "../slot/Slot";

/* -------------------------------------------------------------------------- */
/*                                 FocusBoundary                                 */
/* -------------------------------------------------------------------------- */
interface FocusBoundaryProps extends React.HTMLAttributes<HTMLDivElement> {
  /**
   * FocusBoundary expects a single child element since its a slotted component.
   */
  children: React.ReactElement;
  /**
   * When `true`, tabbing from last item will focus first tabbable
   * and shift+tab from first item will focus last tabbable element.
   * This does not "trap" focus inside the boundary, it only loops it when
   * tabbing. If focus is moved outside the boundary programmatically or by
   * pointer, it will not be moved back.
   *
   * - Hidden inputs (i.e. `<input type="hidden">`) are not considered tabbable.
   * - Elements that are `display: none` or `visibility: hidden` are not considered tabbable.
   * - Elements with `tabIndex < 0` are not considered tabbable.
   * @defaultValue false
   */
  loop?: boolean;
  /**
   * When `true`, focus cannot escape the focus boundary via keyboard,
   * pointer, or a programmatic focus.
   * @defaultValue false
   */
  trapped?: boolean;
  /**
   * Will try to focus the given element on mount.
   *
   * If not provided, FocusBoundary will try to focus the first
   * tabbable element inside the boundary.
   *
   * Set to `false` to not focus anything.
   */
  initialFocus?:
    | boolean
    | React.MutableRefObject<HTMLElement | null>
    | (() => boolean | HTMLElement | null | undefined);
  /**
   * Will try to focus the given element on unmount.
   *
   * If not provided, FocusBoundary will try to focus the element
   * that was focused before the FocusBoundary mounted.
   *
   * Set to `false` to not focus anything.
   */
  returnFocus?:
    | boolean
    | React.MutableRefObject<HTMLElement | null>
    | (() => boolean | HTMLElement | null | undefined);
  /**
   * Hides all outside content from screen readers when true.
   * @default false
   */
  modal?: boolean;
}

const FocusBoundary = forwardRef<HTMLDivElement, FocusBoundaryProps>(
  (
    {
      loop = false,
      trapped = false,
      initialFocus = true,
      returnFocus = true,
      modal = false,
      ...restProps
    }: FocusBoundaryProps,
    forwardedRef,
  ) => {
    const initialFocusRef = useValueAsRef(initialFocus);
    const returnFocusRef = useValueAsRef(returnFocus);

    const lastFocusedElementRef = useRef<HTMLElement | null>(null);
    const [container, setContainer] = useState<HTMLElement | null>(null);
    const mergedRefs = useMergeRefs(forwardedRef, setContainer);

    const focusBoundary = useRef<FocusBoundaryAPI>({
      paused: false,
      pause() {
        this.paused = true;
      },
      resume() {
        this.paused = false;
      },
    }).current;

    /* Handles trapped state */
    useEffect(() => {
      if (!trapped || !container) {
        return;
      }

      function handleFocusIn(event: FocusEvent) {
        if (focusBoundary.paused || container === null) {
          return;
        }

        const target = event.target as HTMLElement | null;
        if (container.contains(target)) {
          lastFocusedElementRef.current = target;
        } else {
          focusElement(lastFocusedElementRef.current, { select: true });
        }
      }

      function handleFocusOut(event: FocusEvent) {
        if (focusBoundary.paused || container === null) {
          return;
        }

        const relatedTarget = event.relatedTarget as HTMLElement | null;
        /*
         * `focusout` event with a `null` `relatedTarget` will happen in a few known cases:
         * 1. When the user switches app/tabs/windows/the browser itself loses focus.
         * 2. In Google Chrome, when the focused element is removed from the DOM.
         * 3. When clicking on an element that cannot receive focus.
         *
         * We let the browser do its thing here because:
         * 1. The browser already keeps a memory of what's focused for when the page gets refocused.
         * 2. In Google Chrome, if we try to focus the deleted focused element (as per below), it
         * throws the CPU to 100%, so we avoid doing anything for this reason here too.
         */
        if (relatedTarget === null) {
          return;
        }

        /*
         * If the focus has moved to an element outside the container, we move focus to the last valid focused element inside.
         * This makes sure to "trap" focus inside the container.
         * We handle focus on focusout instead of focusin to avoid elements recieving focusin events
         * when they are not supposed to (like when clicking on elements outside the container
         */
        if (!container.contains(relatedTarget)) {
          focusElement(lastFocusedElementRef.current, { select: true });
        }
      }

      /**
       * When the currently focused element is removed from the DOM, browsers move focus
       * to the document.body. In this case, we move focus to the container
       * to keep focus trapped correctly instead.
       */
      const handleMutations = (mutations: MutationRecord[]) => {
        if (document.activeElement !== document.body) {
          return;
        }

        if (mutations.some((mutation) => mutation.removedNodes.length > 0)) {
          focusElement(container);
        }
      };

      document.addEventListener("focusin", handleFocusIn);
      document.addEventListener("focusout", handleFocusOut);
      const observer = new MutationObserver(handleMutations);
      observer.observe(container, { childList: true, subtree: true });

      return () => {
        document.removeEventListener("focusin", handleFocusIn);
        document.removeEventListener("focusout", handleFocusOut);
        observer.disconnect();
      };
    }, [trapped, container, focusBoundary.paused]);

    /* Adds element to focus-stack */
    useEffect(() => {
      if (!container) {
        return;
      }

      const ownerDoc = ownerDocument(container);
      const activeElement = ownerDoc.activeElement;
      const closestContainer = activeElement?.closest("[data-focus-boundary]");

      addPreviouslyFocusedElement(ownerDoc.activeElement, closestContainer);
      focusBoundarysStack.add(focusBoundary);

      return () => {
        setTimeout(() => {
          focusBoundarysStack.remove(focusBoundary);
        }, 0);
      };
    }, [container, focusBoundary]);

    /**
     * On unmount, we need to clean up previously focused elements associated with this container
     * This makes sure we don't accidentally try to focus elements that are no longer relevant
     * or will be removed from the DOM.
     */
    useEffect(() => {
      return () => {
        container && deleteContainerAndPreviouslyFocusedElements(container);
      };
    }, [container]);

    useEffect(() => {
      if (!container || !modal) {
        return;
      }

      return hideNonTargetElements([container]);
    }, [container, modal]);

    /* Handles mount focus */
    useClientLayoutEffect(() => {
      if (!container || initialFocusRef.current === false) {
        return;
      }

      const ownerDoc = ownerDocument(container);
      const previouslyFocusedElement = ownerDoc.activeElement;

      queueMicrotask(() => {
        const focusableElements = getTabbableCandidates(container);
        const initialFocusValueOrFn = initialFocusRef.current;
        const resolvedInitialFocus =
          typeof initialFocusValueOrFn === "function"
            ? initialFocusValueOrFn()
            : initialFocusValueOrFn;

        if (
          resolvedInitialFocus === undefined ||
          resolvedInitialFocus === false
        ) {
          return;
        }

        let elToFocus: HTMLElement | null | undefined;
        const fallbackelements = focusableElements[0] || container;

        /* `null` should fallback to default behavior in case of an empty ref. */
        if (resolvedInitialFocus === true || resolvedInitialFocus === null) {
          elToFocus = fallbackelements;
        } else {
          elToFocus = resolveRef(resolvedInitialFocus) || fallbackelements;
        }

        const focusAlreadyInsideFloatingEl = container.contains(
          previouslyFocusedElement,
        );

        if (focusAlreadyInsideFloatingEl) {
          return;
        }

        focusElement(elToFocus, {
          preventScroll: elToFocus === container,
          sync: false,
        });
      });
    }, [container, initialFocusRef]);

    /* Handles unmount focus */
    useClientLayoutEffect(() => {
      if (!container) {
        return;
      }
      const ownerDoc = ownerDocument(container);

      function getReturnElement() {
        const resolvedReturnFocusValueOrFn = returnFocusRef.current;

        let resolvedReturnFocusValue =
          typeof resolvedReturnFocusValueOrFn === "function"
            ? resolvedReturnFocusValueOrFn()
            : resolvedReturnFocusValueOrFn;

        if (
          resolvedReturnFocusValue === undefined ||
          resolvedReturnFocusValue === false
        ) {
          return null;
        }

        /* `null` should fallback to default behavior in case of an empty ref. */
        if (resolvedReturnFocusValue === null) {
          resolvedReturnFocusValue = true;
        }

        if (typeof resolvedReturnFocusValue === "boolean") {
          const el = getPreviouslyFocusedElement();
          return el?.isConnected ? el : ownerDoc.body;
        }

        const fallback = getPreviouslyFocusedElement() || ownerDoc.body;

        return resolveRef(resolvedReturnFocusValue) || fallback;
      }

      return () => {
        const returnElement = getReturnElement() as HTMLElement | null;
        const activeEl = ownerDoc.activeElement;

        queueMicrotask(() => {
          if (
            // eslint-disable-next-line react-hooks/exhaustive-deps
            returnFocusRef.current &&
            returnElement &&
            returnElement !== activeEl
          ) {
            returnElement.focus({ preventScroll: true });
          }
        });
      };
    }, [container, returnFocusRef]);

    /* Takes care of looping focus */
    const handleKeyDown = useCallback(
      (event: React.KeyboardEvent) => {
        if ((!loop && !trapped) || focusBoundary.paused) {
          return;
        }

        const isTabKey =
          event.key === "Tab" &&
          !event.altKey &&
          !event.ctrlKey &&
          !event.metaKey;

        const focusedElement = document.activeElement;

        if (isTabKey && focusedElement) {
          const containerTarget = event.currentTarget as HTMLElement;
          const [first, last] = getTabbableEdges(containerTarget);

          /* We can only wrap focus if we have tabbable edges */
          if (!(first && last)) {
            /*
             * No need to do anything if active element is the expected focus-target
             * Case: No tabbable elements, focus should stay on container. If we don't preventDefault, the container will lose focus
             * and potentially lose controll of focus to browser (like focusing address bar).
             */
            if (focusedElement === containerTarget) {
              event.preventDefault();
            }
            return;
          }

          /**
           * Since we are either trapped + looping, or one of them we will do nothing when trapped and focus first element when looping.
           */
          if (!event.shiftKey && focusedElement === last) {
            event.preventDefault();
            if (loop) {
              focusElement(first, { select: true });
            }
          } else if (event.shiftKey && focusedElement === first) {
            event.preventDefault();
            if (loop) {
              focusElement(last, { select: true });
            }
          }
        }
      },
      [loop, trapped, focusBoundary.paused],
    );

    return (
      <Slot
        tabIndex={-1}
        {...restProps}
        ref={mergedRefs}
        onKeyDown={handleKeyDown}
        data-focus-boundary
      />
    );
  },
);

/* ---------------------------- FocusBoundary utils ---------------------------- */
/**
 * Returns the first and last tabbable elements inside a container as a tuple.
 */
function getTabbableEdges(container: HTMLElement) {
  const candidates = getTabbableCandidates(container, { omitLinks: false });
  return [
    findFirstVisible(candidates, container),
    findFirstVisible(candidates.reverse(), container),
  ] as const;
}

/**
 * Returns the first visible element in a list.
 * NOTE: Only checks visibility up to the `container`.
 */
function findFirstVisible(elements: HTMLElement[], container: HTMLElement) {
  for (const element of elements) {
    if (!isHidden(element, { upTo: container })) {
      return element;
    }
  }
}

function isHidden(node: HTMLElement, { upTo }: { upTo?: HTMLElement }) {
  if (getComputedStyle(node).visibility === "hidden") {
    return true;
  }

  while (node) {
    /* we stop at `upTo` */
    if (upTo !== undefined && node === upTo) {
      return false;
    }
    if (getComputedStyle(node).display === "none") {
      return true;
    }
    node = node.parentElement as HTMLElement;
  }
  return false;
}

/* ---------------------------- FocusBoundary stack ---------------------------- */
type FocusBoundaryAPI = { paused: boolean; pause(): void; resume(): void };
const focusBoundarysStack = createFocusBoundarysStack();

function createFocusBoundarysStack() {
  /* A stack of focus-boundaries, with the active one at the top */
  let stack: FocusBoundaryAPI[] = [];

  return {
    add(focusBoundary: FocusBoundaryAPI) {
      /* Pause the currently active focus-boundary (at the top of the stack) */
      const activeFocusBoundary = stack[0];
      if (focusBoundary !== activeFocusBoundary) {
        activeFocusBoundary?.pause();
      }
      /* remove in case it already exists (because we'll re-add it at the top of the stack) */
      stack = arrayRemove(stack, focusBoundary);
      stack.unshift(focusBoundary);
    },

    remove(focusBoundary: FocusBoundaryAPI) {
      stack = arrayRemove(stack, focusBoundary);
      stack[0]?.resume();
    },
  };
}

function arrayRemove<T>(array: T[], item: T) {
  const updatedArray = [...array];
  const index = updatedArray.indexOf(item);
  if (index !== -1) {
    updatedArray.splice(index, 1);
  }
  return updatedArray;
}

const LIST_LIMIT = 10;
let previouslyFocusedElements: Element[] = [];
const focusedElementsByContainer = new WeakMap<Element, Element[]>();

function clearDisconnectedPreviouslyFocusedElements() {
  previouslyFocusedElements = previouslyFocusedElements.filter(
    (el) => el.isConnected,
  );
}

/**
 * Removes "will be" unmounted elements from previouslyFocusedElements,
 * and deletes the container from focusedElementsByContainer.
 */
function deleteContainerAndPreviouslyFocusedElements(container: HTMLElement) {
  const nestedElements = focusedElementsByContainer.get(container) || [];
  previouslyFocusedElements = previouslyFocusedElements.filter((el) => {
    return !nestedElements.includes(el);
  });
  focusedElementsByContainer.delete(container);
}

function addPreviouslyFocusedElement(
  element: Element | null,
  container: Element | null | undefined,
) {
  clearDisconnectedPreviouslyFocusedElements();
  if (element && element?.nodeName !== "BODY") {
    previouslyFocusedElements.push(element);

    if (container) {
      const nestedElements = focusedElementsByContainer.get(container) || [];
      nestedElements.push(element);
      focusedElementsByContainer.set(container, nestedElements);
    }

    if (previouslyFocusedElements.length > LIST_LIMIT) {
      previouslyFocusedElements = previouslyFocusedElements.slice(-LIST_LIMIT);
    }
  }
}

function getPreviouslyFocusedElement() {
  clearDisconnectedPreviouslyFocusedElements();
  return previouslyFocusedElements[previouslyFocusedElements.length - 1];
}

export { FocusBoundary };
export type { FocusBoundaryProps };