"use client"

import * as React from "react"

import { cn } from "@/shared/lib/utils"

export interface FocusTrapProps {
  children: React.ReactNode
  /** 是否激活焦点陷阱，默认 true */
  active?: boolean
  /** 初始焦点元素，未指定时聚焦第一个可聚焦元素 */
  initialFocusRef?: React.RefObject<HTMLElement | null>
  /** 关闭时是否恢复焦点到触发元素，默认 true */
  restoreFocus?: boolean
  className?: string
}

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

/**
 * 焦点陷阱组件（用于模态框/对话框）。
 * - 捕获 Tab/Shift+Tab 在容器内循环
 * - 支持初始焦点元素
 * - 支持恢复焦点到触发元素
 */
export function FocusTrap({
  children,
  active = true,
  initialFocusRef,
  restoreFocus = true,
  className,
}: FocusTrapProps) {
  const containerRef = React.useRef<HTMLDivElement>(null)
  const previouslyFocusedRef = React.useRef<HTMLElement | null>(null)

  React.useEffect(() => {
    if (!active) return
    previouslyFocusedRef.current = document.activeElement as HTMLElement | null

    const container = containerRef.current
    if (container) {
      const focusTarget =
        initialFocusRef?.current ?? getFirstFocusable(container)
      if (focusTarget) {
        focusTarget.focus()
      }
    }

    return () => {
      if (restoreFocus && previouslyFocusedRef.current) {
        previouslyFocusedRef.current.focus()
      }
    }
  }, [active, initialFocusRef, restoreFocus])

  const handleKeyDown = React.useCallback(
    (event: React.KeyboardEvent<HTMLDivElement>) => {
      if (event.key !== "Tab") return
      const container = containerRef.current
      if (!container) return

      const focusables = getFocusables(container)
      if (focusables.length === 0) {
        event.preventDefault()
        container.focus()
        return
      }

      const first = focusables[0]
      const last = focusables[focusables.length - 1]
      const activeEl = document.activeElement

      if (event.shiftKey) {
        if (activeEl === first || !container.contains(activeEl)) {
          event.preventDefault()
          last.focus()
        }
      } else {
        if (activeEl === last || !container.contains(activeEl)) {
          event.preventDefault()
          first.focus()
        }
      }
    },
    []
  )

  if (!active) {
    return <>{children}</>
  }

  return (
    <div
      ref={containerRef}
      onKeyDown={handleKeyDown}
      tabIndex={-1}
      className={cn("outline-none", className)}
    >
      {children}
    </div>
  )
}

function getFocusables(container: HTMLElement): HTMLElement[] {
  return Array.from(
    container.querySelectorAll<HTMLElement>(FOCUSABLE_SELECTOR)
  ).filter((el) => {
    if (el.hasAttribute("disabled")) return false
    if (el.getAttribute("aria-hidden") === "true") return false
    return el.offsetParent !== null || el.getClientRects().length > 0
  })
}

function getFirstFocusable(container: HTMLElement): HTMLElement | null {
  return getFocusables(container)[0] ?? null
}
