import { useEffect, useLayoutEffect, useRef } from 'react'
import type { ReactNode, RefObject } from 'react'
import { focusableElements, topmostDialog } from './focusable'
import './ModalShell.css'

export interface ModalShellProps {
  isOpen: boolean
  onClose: () => void
  titleId: string
  children: ReactNode
  contentClassName?: string
  overlayClassName?: string
  backdropClassName?: string
  initialFocusRef?: RefObject<HTMLElement>
}

export function ModalShell({
  isOpen,
  onClose,
  titleId,
  children,
  contentClassName,
  overlayClassName,
  backdropClassName,
  initialFocusRef,
}: ModalShellProps) {
  const contentRef = useRef<HTMLDivElement>(null)
  const onCloseRef = useRef(onClose)
  const returnFocusRef = useRef<HTMLElement | null>(null)

  useEffect(() => { onCloseRef.current = onClose }, [onClose])

  useEffect(() => {
    if (!isOpen || !contentRef.current) return
    const content = contentRef.current
    const syncTopmost = () => {
      const topmost = topmostDialog()
      const isTopmost = topmost === content
      content.setAttribute('aria-modal', String(isTopmost))
      // Dialogs rendered in a portal above us should hide the underlying
      // dialog from assistive technology until the topmost one closes.
      if (!isTopmost && topmost && !content.contains(topmost)) content.setAttribute('aria-hidden', 'true')
      else content.removeAttribute('aria-hidden')
    }
    syncTopmost()
    const observer = new MutationObserver(syncTopmost)
    observer.observe(document.body, { childList: true, subtree: true })
    return () => observer.disconnect()
  }, [isOpen])

  useLayoutEffect(() => {
    if (!isOpen) return

    returnFocusRef.current = document.activeElement instanceof HTMLElement
      && document.activeElement !== document.body
      ? document.activeElement
      : null
    if (document.activeElement instanceof HTMLElement && document.activeElement !== document.body) {
      document.activeElement.blur()
    }
    if (initialFocusRef?.current) initialFocusRef.current.focus()
    else contentRef.current?.focus()

    function handleKeyDown(event: KeyboardEvent) {
      if (event.defaultPrevented) return
      if (topmostDialog() !== contentRef.current) return
      if (event.key === 'Escape') {
        event.preventDefault()
        onCloseRef.current()
        return
      }

      if (event.key !== 'Tab' || !contentRef.current) return
      const focusable = focusableElements(contentRef.current)
      if (focusable.length === 0) {
        event.preventDefault()
        contentRef.current.focus()
        return
      }

      const first = focusable[0]
      const last = focusable[focusable.length - 1]
      if (event.shiftKey && (document.activeElement === contentRef.current || document.activeElement === first || !contentRef.current.contains(document.activeElement))) {
        event.preventDefault()
        last.focus()
      } else if (!event.shiftKey && (document.activeElement === last || !contentRef.current.contains(document.activeElement))) {
        event.preventDefault()
        first.focus()
      }
    }

    function handleFocusIn(event: FocusEvent) {
      const content = contentRef.current
      if (!content || topmostDialog() !== content || content.contains(event.target as Node)) return
      const preferred = initialFocusRef?.current
      if (preferred?.isConnected && content.contains(preferred)) preferred.focus({ preventScroll: true })
      else content.focus({ preventScroll: true })
    }

    document.addEventListener('keydown', handleKeyDown)
    document.addEventListener('focusin', handleFocusIn, true)
    return () => {
      document.removeEventListener('keydown', handleKeyDown)
      document.removeEventListener('focusin', handleFocusIn, true)
      const top = topmostDialog()
      const anotherDialogIsOnTop = top !== undefined && top !== contentRef.current
      if (!anotherDialogIsOnTop && (document.activeElement === document.body || contentRef.current?.contains(document.activeElement)) && returnFocusRef.current?.isConnected) {
        returnFocusRef.current.focus()
      }
    }
  }, [isOpen, initialFocusRef])

  if (!isOpen) return null

  return (
    <div className={`modal-shell-overlay${overlayClassName ? ` ${overlayClassName}` : ''}`}>
      <div className={`modal-shell-backdrop${backdropClassName ? ` ${backdropClassName}` : ''}`} aria-hidden="true" onClick={() => onCloseRef.current()} />
      <div
        ref={contentRef}
        className={`modal-shell-content${contentClassName ? ` ${contentClassName}` : ''}`}
        role="dialog"
        aria-modal="true"
        aria-labelledby={titleId}
        tabIndex={-1}
      >
        {children}
      </div>
    </div>
  )
}