// SPDX-License-Identifier: MIT
import { useEffect, useRef } from "react"

interface FocusTrapProps {
  container?: React.RefObject<HTMLElement> | null
  initialFocusRef?: React.RefObject<HTMLElement>
}

export function FocusTrap({ container, initialFocusRef }: FocusTrapProps): null {
  const refLastFocus = useRef<HTMLElement | null>()
  /**
   * Handle focus lock on the modal
   */
  useEffect(() => {
    const handleKeyEvent = (event: KeyboardEvent) => {
      if (container?.current) {
        tabTrappingKey(event, container.current)
      }
    }

    document.addEventListener("keydown", handleKeyEvent)
    // On mount we focus on the first focusable element in the modal if there is one
    if (container?.current) {
      const savePreviousFocus = () => {
        // First we save the last focused element
        // only if it's a focusable element
        if (
          candidateSelectors.findIndex(
            selector => document.activeElement?.matches(selector),
          ) !== -1
        ) {
          refLastFocus.current = document.activeElement as HTMLElement
        }
      }

      if (initialFocusRef) {
        savePreviousFocus()
        // We need to schedule focusing on a next frame - this allows to focus on the modal root
        requestAnimationFrame(() => {
          initialFocusRef.current?.focus()
        })
      } else {
        const allTabbingElements = getAllTabbingElements(container.current)
        if (allTabbingElements[0]) {
          savePreviousFocus()
          allTabbingElements[0].focus()
        }
      }
    }
    return () => {
      document.removeEventListener("keydown", handleKeyEvent)
      // On unmount we restore the focus to the last focused element
      refLastFocus.current?.focus()
    }
  }, [container, initialFocusRef])

  return null
}

// https://github.com/alexandrzavalii/focus-trap-js/blob/master/src/index.js v1.1.0
const candidateSelectors = [
  "input",
  "select",
  "textarea",
  "a[href]",
  "button",
  "[tabindex]",
  "audio[controls]",
  "video[controls]",
  '[contenteditable]:not([contenteditable="false"])',
]

function isInput(node: HTMLElement): node is HTMLInputElement {
  return node.tagName === "INPUT"
}

function isNotRadioOrTabbableRadio(node: HTMLElement) {
  if (!isInput(node) || node.type !== "radio" || !node.name) {
    return true
  }
  const radioSet = (node.form ?? node.ownerDocument).querySelectorAll<HTMLInputElement>(
    `input[type="radio"][name="${node.name}"]`,
  )
  const checked = Array.from(radioSet).find(
    radio => radio.checked && radio.form === node.form,
  )
  return checked === node || (checked === undefined && radioSet[0] === node)
}

function getAllTabbingElements(parentElem: HTMLElement) {
  const { activeElement } = document
  return Array.from<HTMLElement>(
    parentElem.querySelectorAll(candidateSelectors.join(",")),
  ).filter(
    node =>
      activeElement === node ||
      (!(node as HTMLElement & { disabled?: boolean }).disabled &&
        getTabindex(node) > -1 &&
        // check if hidden: offsetParent being null will allow detecting cases where an element is invisible or inside an invisible element,
        // as long as the element does not use position: fixed. For them, their visibility has to be checked directly as well.
        !(node.offsetParent === null || getComputedStyle(node).visibility === "hidden") &&
        isNotRadioOrTabbableRadio(node)),
  )
}

function tabTrappingKey(event: KeyboardEvent, parentElem: HTMLElement) {
  // check if current event keyCode is tab
  if (!event || event.key !== "Tab") return

  if (!parentElem.contains(event.target as HTMLElement)) {
    return false
  }

  const allTabbingElements = getAllTabbingElements(parentElem)
  const [firstFocusableElement] = allTabbingElements
  const lastFocusableElement = allTabbingElements.at(-1)!

  if (event.shiftKey && event.target === firstFocusableElement) {
    lastFocusableElement.focus()
    event.preventDefault()
    return true
  } else if (!event.shiftKey && event.target === lastFocusableElement) {
    firstFocusableElement.focus()
    event.preventDefault()
    return true
  }
  return false
}

function getTabindex(node: HTMLElement) {
  const tabindexAttr = parseInt(node.getAttribute("tabindex")!, 10)

  if (!isNaN(tabindexAttr)) return tabindexAttr
  // Browsers do not return tabIndex correctly for contentEditable nodes;
  // so if they don't have a tabindex attribute specifically set, assume it's 0.

  if (node.hasAttribute("contentEditable")) return 0
  return node.tabIndex
}
