/**
 * Focus management helpers for focusable and tabbable elements.
 * @module Focus utilities
 */

import { contains, getActiveElement, isFrame, isVisible } from "./dom.ts";

const selector =
  "input:not([type='hidden']):not([disabled]), select:not([disabled]), " +
  "textarea:not([disabled]), a[href], button:not([disabled]), [tabindex], " +
  "summary, iframe, object, embed, area[href], audio[controls], " +
  "video[controls], [contenteditable]:not([contenteditable='false'])";

function hasNegativeTabIndex(element: Element) {
  const tabIndex = Number.parseInt(element.getAttribute("tabindex") || "0", 10);
  return tabIndex < 0;
}

/**
 * Checks whether `element` is focusable or not.
 * @example
 * isFocusable(document.querySelector("input")); // true
 * isFocusable(document.querySelector("input[tabindex='-1']")); // true
 * isFocusable(document.querySelector("input[hidden]")); // false
 * isFocusable(document.querySelector("input:disabled")); // false
 */
export function isFocusable(element: Element) {
  if (!element.matches(selector)) return false;
  if (!isVisible(element)) return false;
  if (element.closest("[inert]")) return false;
  return true;
}

/**
 * Checks whether `element` is tabbable or not.
 * @example
 * isTabbable(document.querySelector("input")); // true
 * isTabbable(document.querySelector("input[tabindex='-1']")); // false
 * isTabbable(document.querySelector("input[hidden]")); // false
 * isTabbable(document.querySelector("input:disabled")); // false
 */
export function isTabbable(
  element: Element | HTMLElement | HTMLInputElement,
): element is HTMLElement {
  if (!isFocusable(element)) return false;
  if (hasNegativeTabIndex(element)) return false;
  // If the element is a radio button in a form, we must take roving tabindex
  // into account.
  if (!("form" in element)) return true;
  if (!element.form) return true;
  if (element.checked) return true;
  if (element.type !== "radio") return true;
  // If the radio button is not part of a radio group, it's tabbable.
  const radioGroup = element.form.elements.namedItem(element.name);
  if (!radioGroup) return true;
  if (!("length" in radioGroup)) return true;
  // If we are in a radio group, we must check if the active element is part of
  // the same group, which would make all the other radio buttons in the group
  // non-tabbable.
  const activeElement = getActiveElement(element);
  if (!activeElement) return true;
  if (activeElement === element) return true;
  if (!("form" in activeElement)) return true;
  if (activeElement.form !== element.form) return true;
  if (!("name" in activeElement)) return true;
  if (activeElement.name !== element.name) return true;
  return false;
}

/**
 * Returns all the focusable elements in `container`.
 */
export function getAllFocusableIn(
  container: HTMLElement,
  includeContainer?: boolean,
) {
  const elements = Array.from(
    container.querySelectorAll<HTMLElement>(selector),
  );
  if (includeContainer) {
    elements.unshift(container);
  }

  const focusableElements = elements.filter(isFocusable);

  focusableElements.forEach((element, i) => {
    if (!isFrame(element)) return;
    const frameBody = element.contentDocument?.body;
    if (!frameBody) return;
    focusableElements.splice(i, 1, ...getAllFocusableIn(frameBody));
  });

  return focusableElements;
}

/**
 * Returns all the focusable elements in the document.
 */
export function getAllFocusable(includeBody?: boolean) {
  return getAllFocusableIn(document.body, includeBody);
}

/**
 * Returns the first focusable element in `container`.
 */
export function getFirstFocusableIn(
  container: HTMLElement,
  includeContainer?: boolean,
) {
  const [first] = getAllFocusableIn(container, includeContainer);
  return first || null;
}

/**
 * Returns the first focusable element in the document.
 */
export function getFirstFocusable(includeBody?: boolean) {
  return getFirstFocusableIn(document.body, includeBody);
}

/**
 * Returns all the tabbable elements in `container`, including the container
 * itself.
 */
export function getAllTabbableIn(
  container: HTMLElement,
  includeContainer?: boolean,
  fallbackToFocusable?: boolean,
) {
  const elements = Array.from(
    container.querySelectorAll<HTMLElement>(selector),
  );
  const tabbableElements = elements.filter(isTabbable);

  if (includeContainer && isTabbable(container)) {
    tabbableElements.unshift(container);
  }

  tabbableElements.forEach((element, i) => {
    if (!isFrame(element)) return;
    const frameBody = element.contentDocument?.body;
    if (!frameBody) return;
    const allFrameTabbable = getAllTabbableIn(
      frameBody,
      false,
      fallbackToFocusable,
    );
    tabbableElements.splice(i, 1, ...allFrameTabbable);
  });

  if (!tabbableElements.length && fallbackToFocusable) {
    // Filter the fallback to focusable elements (and expand iframes like the
    // tabbable path) so callers don't get non-focusable matches such as a
    // `display: none` input, which a `focus()` call silently ignores.
    return getAllFocusableIn(container);
  }
  return tabbableElements;
}

/**
 * Returns all the tabbable elements in the document.
 */
export function getAllTabbable(fallbackToFocusable?: boolean) {
  return getAllTabbableIn(document.body, false, fallbackToFocusable);
}

/**
 * Returns the first tabbable element in `container`, including the container
 * itself if it's tabbable.
 */
export function getFirstTabbableIn(
  container: HTMLElement,
  includeContainer?: boolean,
  fallbackToFocusable?: boolean,
): HTMLElement | null {
  // Unlike getAllTabbableIn, this returns as soon as a tabbable element is
  // found. Each isTabbable check queries style and layout, so collecting every
  // tabbable element just to take the first one is wasteful in containers with
  // many tabbable elements, such as dialogs with forms.
  if (includeContainer && isTabbable(container)) {
    if (!isFrame(container)) return container;
    const containerFrameBody = container.contentDocument?.body;
    if (!containerFrameBody) return container;
    const frameTabbable = getFirstTabbableIn(
      containerFrameBody,
      false,
      fallbackToFocusable,
    );
    if (frameTabbable) return frameTabbable;
    // The container frame has no tabbable content. Fall through to the
    // container's own descendants, like getAllTabbableIn, which removed the
    // empty frame from the list and continued.
  }
  const elements = container.querySelectorAll<HTMLElement>(selector);
  for (const element of elements) {
    if (!isTabbable(element)) continue;
    if (isFrame(element)) {
      const frameBody = element.contentDocument?.body;
      if (!frameBody) return element;
      const frameTabbable = getFirstTabbableIn(
        frameBody,
        false,
        fallbackToFocusable,
      );
      if (frameTabbable) return frameTabbable;
      continue;
    }
    return element;
  }
  if (fallbackToFocusable) {
    // Filter the fallback to focusable elements so callers don't get a
    // non-focusable match such as a `display: none` input, which a `focus()`
    // call silently ignores.
    return getFirstFocusableIn(container);
  }
  return null;
}

/**
 * Returns the first tabbable element in the document.
 */
export function getFirstTabbable(fallbackToFocusable?: boolean) {
  return getFirstTabbableIn(document.body, false, fallbackToFocusable);
}

/**
 * Returns the last tabbable element in `container`, including the container
 * itself if it's tabbable.
 */
export function getLastTabbableIn(
  container: HTMLElement,
  includeContainer?: boolean,
  fallbackToFocusable?: boolean,
) {
  const allTabbable = getAllTabbableIn(
    container,
    includeContainer,
    fallbackToFocusable,
  );
  return allTabbable[allTabbable.length - 1] || null;
}

/**
 * Returns the last tabbable element in the document.
 */
export function getLastTabbable(fallbackToFocusable?: boolean) {
  return getLastTabbableIn(document.body, false, fallbackToFocusable);
}

interface GetTabbableInDirectionParams {
  container: HTMLElement;
  includeContainer?: boolean;
  reverse?: boolean;
  fallbackToEdge?: boolean;
  fallbackToFocusable?: boolean;
}

function getTabbableInDirection({
  container,
  includeContainer,
  reverse,
  fallbackToEdge,
  fallbackToFocusable,
}: GetTabbableInDirectionParams) {
  const activeElement = getActiveElement(container);
  const allFocusable = getAllFocusableIn(container, includeContainer);
  if (reverse) {
    allFocusable.reverse();
  }
  const activeIndex = allFocusable.indexOf(activeElement as HTMLElement);
  const candidates = allFocusable.slice(activeIndex + 1);
  return (
    candidates.find(isTabbable) ||
    (fallbackToEdge ? allFocusable.find(isTabbable) : null) ||
    (fallbackToFocusable ? candidates[0] : null) ||
    null
  );
}

/**
 * Returns the next tabbable element in `container`.
 */
export function getNextTabbableIn(
  container: HTMLElement,
  includeContainer?: boolean,
  fallbackToFirst?: boolean,
  fallbackToFocusable?: boolean,
) {
  return getTabbableInDirection({
    container,
    includeContainer,
    fallbackToEdge: fallbackToFirst,
    fallbackToFocusable,
  });
}

/**
 * Returns the next tabbable element in the document.
 */
export function getNextTabbable(
  fallbackToFirst?: boolean,
  fallbackToFocusable?: boolean,
) {
  return getNextTabbableIn(
    document.body,
    false,
    fallbackToFirst,
    fallbackToFocusable,
  );
}

/**
 * Returns the previous tabbable element in `container`.
 *
 */
export function getPreviousTabbableIn(
  container: HTMLElement,
  includeContainer?: boolean,
  fallbackToLast?: boolean,
  fallbackToFocusable?: boolean,
) {
  return getTabbableInDirection({
    container,
    includeContainer,
    reverse: true,
    fallbackToEdge: fallbackToLast,
    fallbackToFocusable,
  });
}

/**
 * Returns the previous tabbable element in the document.
 */
export function getPreviousTabbable(
  fallbackToLast?: boolean,
  fallbackToFocusable?: boolean,
) {
  return getPreviousTabbableIn(
    document.body,
    false,
    fallbackToLast,
    fallbackToFocusable,
  );
}

/**
 * Returns the closest focusable element.
 */
export function getClosestFocusable(element?: HTMLElement | null) {
  // Start at the parent so a selector-matching but non-focusable current
  // element cannot be returned repeatedly and loop forever.
  while (element && !isFocusable(element)) {
    element = element.parentElement?.closest<HTMLElement>(selector) ?? null;
  }
  return element || null;
}

/**
 * Checks if `element` has focus. Elements that are referenced by
 * `aria-activedescendant` are also considered.
 * @example
 * hasFocus(document.getElementById("id"));
 */
export function hasFocus(element: Element) {
  const activeElement = getActiveElement(element);
  if (!activeElement) return false;
  if (activeElement === element) return true;
  const activeDescendant = activeElement.getAttribute("aria-activedescendant");
  if (!activeDescendant) return false;
  return activeDescendant === element.id;
}

/**
 * Checks if `element` has focus within. Elements that are referenced by
 * `aria-activedescendant` are also considered.
 * @example
 * hasFocusWithin(document.getElementById("id"));
 */
export function hasFocusWithin(element: Node | Element) {
  const activeElement = getActiveElement(element);
  if (!activeElement) return false;
  if (contains(element, activeElement)) return true;
  const activeDescendant = activeElement.getAttribute("aria-activedescendant");
  if (!activeDescendant) return false;
  if (!("id" in element)) return false;
  if (activeDescendant === element.id) return true;
  return !!element.querySelector(`#${CSS.escape(activeDescendant)}`);
}

/**
 * Focus on an element only if it's not already focused.
 */
export function focusIfNeeded(element: HTMLElement) {
  if (!hasFocusWithin(element) && isFocusable(element)) {
    element.focus();
  }
}

/**
 * Disable focus on `element`.
 */
export function disableFocus(element: HTMLElement) {
  const currentTabindex = element.getAttribute("tabindex") ?? "";
  element.setAttribute("data-tabindex", currentTabindex);
  element.setAttribute("tabindex", "-1");
}

/**
 * Makes elements inside container not tabbable.
 */
export function disableFocusIn(
  container: HTMLElement,
  includeContainer?: boolean,
) {
  const tabbableElements = getAllTabbableIn(container, includeContainer);
  for (const element of tabbableElements) {
    disableFocus(element);
  }
}

/**
 * Restores tabbable elements inside container that were affected by
 * disableFocusIn.
 */
export function restoreFocusIn(container: HTMLElement) {
  const elements = container.querySelectorAll<HTMLElement>("[data-tabindex]");
  const restoreTabIndex = (element: HTMLElement) => {
    const tabindex = element.getAttribute("data-tabindex");
    element.removeAttribute("data-tabindex");
    if (tabindex) {
      element.setAttribute("tabindex", tabindex);
    } else {
      element.removeAttribute("tabindex");
    }
  };
  if (container.hasAttribute("data-tabindex")) {
    restoreTabIndex(container);
  }
  for (const element of elements) {
    restoreTabIndex(element);
  }
}

/**
 * Focus on element and scroll into view.
 */
export function focusIntoView(
  element: HTMLElement,
  options?: ScrollIntoViewOptions,
) {
  if (!("scrollIntoView" in element)) {
    // @ts-expect-error
    element.focus();
  } else {
    element.focus({ preventScroll: true });
    element.scrollIntoView({ block: "nearest", inline: "nearest", ...options });
  }
}