import type { Ref } from 'vue'
import { nextTick, onUnmounted, watch } from 'vue'

const FOCUSABLE_SELECTOR = [
  'a[href]',
  'button:not([disabled])',
  'input:not([disabled])',
  'select:not([disabled])',
  'textarea:not([disabled])',
  '[tabindex]:not([tabindex="-1"])',
].join(',')

function getFocusable(container: HTMLElement): HTMLElement[] {
  return Array.from(container.querySelectorAll<HTMLElement>(FOCUSABLE_SELECTOR))
}

export interface FocusTrapOptions {
  initialFocus?: () => HTMLElement | null
}

export function useFocusTrap(
  containerRef: Ref<HTMLElement | null>,
  active: Ref<boolean>,
  options: FocusTrapOptions = {},
) {
  let previouslyFocused: HTMLElement | null = null

  function handleKeydown(event: KeyboardEvent) {
    if (event.key !== 'Tab') return
    const container = containerRef.value
    if (!container) return

    const focusable = getFocusable(container)
    if (focusable.length === 0) {
      event.preventDefault()
      return
    }

    const first = focusable[0]!
    const last = focusable[focusable.length - 1]!
    const current = document.activeElement

    if (event.shiftKey) {
      if (current === first || !container.contains(current)) {
        event.preventDefault()
        last.focus()
      }
    } else {
      if (current === last || !container.contains(current)) {
        event.preventDefault()
        first.focus()
      }
    }
  }

  watch(
    active,
    (isActive) => {
      if (isActive) {
        previouslyFocused = document.activeElement as HTMLElement | null
        document.addEventListener('keydown', handleKeydown)
        nextTick(() => {
          const container = containerRef.value
          if (!container) return
          if (container.contains(document.activeElement)) return
          const target = options.initialFocus?.() ?? getFocusable(container)[0] ?? container
          target?.focus()
        })
      } else {
        document.removeEventListener('keydown', handleKeydown)
        previouslyFocused?.focus?.()
        previouslyFocused = null
      }
    },
    { immediate: true },
  )

  onUnmounted(() => {
    document.removeEventListener('keydown', handleKeydown)
  })
}