import * as React from "react"
import * as DialogPrimitive from "@radix-ui/react-dialog"
import { X } from "lucide-react"

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

const Dialog = DialogPrimitive.Root

const DialogTrigger = DialogPrimitive.Trigger

const DialogPortal = DialogPrimitive.Portal

const DialogClose = DialogPrimitive.Close

/**
 * Focus capture and restore, for every overlay in the product.
 *
 * Radix returns focus to whatever opened the overlay only when it can still
 * find it, and on several of PIG's surfaces it cannot: sheets opened
 * programmatically have no `SheetTrigger` at all, and the palette's opener is a
 * control the next route unmounts. Measured, Escape from nine overlays left
 * focus on `<body>` — on `/piggy` that is 121 Tab presses back to where the
 * person was, which is a keyboard trap wearing a dismissal.
 *
 * The pattern was written once at `CommandPalette.tsx` and is hoisted here so
 * every dialog, sheet, drawer and alert inherits it without a call-site change.
 *
 * The capture hangs off the content element's ref rather than the wrapper's
 * first render, and that distinction is the whole component. `SheetContent`
 * renders on every render of the page that declares it — Radix gates on open
 * *below* it, inside the portal — so reading `document.activeElement` while our
 * own function body runs reads it at page load, which is `<body>`, which is the
 * bug being fixed. A ref callback fires only when the content genuinely mounts,
 * and refs are attached earlier in the commit than the effect Radix's
 * FocusScope uses to move focus into the overlay.
 */
export function useOverlayFocusRestore<T extends HTMLElement>(
  forwardedRef: React.ForwardedRef<T>,
  onCloseAutoFocus?: (event: Event) => void,
): {
  ref: (node: T | null) => void
  onCloseAutoFocus: (event: Event) => void
} {
  const opener = React.useRef<HTMLElement | null>(null)
  const captured = React.useRef(false)
  const forwarded = React.useRef(forwardedRef)
  forwarded.current = forwardedRef

  // Deliberately stable: a ref callback that changed identity would be called
  // with null and then the node again mid-open, and the second capture would
  // read a control inside the overlay as the opener.
  const ref = React.useCallback((node: T | null) => {
    if (node) {
      if (!captured.current) {
        captured.current = true
        const active = document.activeElement
        opener.current =
          active instanceof HTMLElement && active !== document.body ? active : null
      }
    } else {
      // Armed for the next open. `opener` itself survives, because the close
      // handler below runs in a passive effect cleanup — after React has
      // already detached this ref.
      captured.current = false
    }
    const target = forwarded.current
    if (typeof target === "function") target(node)
    else if (target) target.current = node
  }, [])

  const handleCloseAutoFocus = React.useCallback(
    (event: Event) => {
      onCloseAutoFocus?.(event)
      // The call site wins. CommandPalette has its own restore for the case
      // where choosing an item navigates away from the control that opened it.
      if (event.defaultPrevented) return
      const target = opener.current
      // `isConnected` because closing may have navigated, leaving the opener
      // detached — focusing a node in no document does nothing.
      //
      // The fallback is the page's own content landmark rather than Radix's,
      // which lands on `<body>`. That case is real and now reachable: an
      // account's contract row navigates to `/contracts?contract=…`, which
      // opens the detail sheet on a page the opener never existed on, so
      // dismissing it dropped the keyboard back to the top of the document.
      // `#page-content` is the same target the skip link uses and already
      // carries `tabIndex={-1}` for exactly this.
      if (!target || !target.isConnected) {
        const content = document.getElementById('page-content')
        if (!content) return
        event.preventDefault()
        content.focus()
        return
      }
      event.preventDefault()
      target.focus()
    },
    [onCloseAutoFocus],
  )

  return { ref, onCloseAutoFocus: handleCloseAutoFocus }
}

const DialogOverlay = React.forwardRef<
  React.ElementRef<typeof DialogPrimitive.Overlay>,
  React.ComponentPropsWithoutRef<typeof DialogPrimitive.Overlay>
>(({ className, ...props }, ref) => (
  <DialogPrimitive.Overlay
    ref={ref}
    className={cn(
      "fixed inset-0 z-50 bg-black/80 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=open]:duration-3 data-[state=closed]:duration-3 data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=open]:ease-enter data-[state=closed]:ease-exit",
      className
    )}
    {...props}
  />
))
DialogOverlay.displayName = DialogPrimitive.Overlay.displayName

const DialogContent = React.forwardRef<
  React.ElementRef<typeof DialogPrimitive.Content>,
  React.ComponentPropsWithoutRef<typeof DialogPrimitive.Content>
>(({ className, children, onCloseAutoFocus, ...props }, ref) => {
  const focus = useOverlayFocusRestore(ref, onCloseAutoFocus)
  return (
    <DialogPortal>
      <DialogOverlay />
      <DialogPrimitive.Content
        ref={focus.ref}
        onCloseAutoFocus={focus.onCloseAutoFocus}
        className={cn(
          "fixed left-[50%] top-[50%] z-50 grid w-[calc(100%-2rem)] max-w-lg translate-x-[-50%] translate-y-[-50%] gap-4 rounded-2xl border border-border bg-surface p-5 shadow-lg data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=open]:duration-3 data-[state=closed]:duration-3 data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[state=closed]:slide-out-to-left-1/2 data-[state=closed]:slide-out-to-top-[48%] data-[state=open]:slide-in-from-left-1/2 data-[state=open]:slide-in-from-top-[48%] data-[state=open]:ease-enter data-[state=closed]:ease-exit",
          className
        )}
        {...props}
      >
        {children}
        <DialogPrimitive.Close className="absolute right-2 top-2 z-10 flex size-11 items-center justify-center rounded-lg text-muted opacity-70 transition-colors duration-1 hover:bg-surface-2 hover:opacity-100 disabled:pointer-events-none">
          <X className="h-4 w-4" />
          <span className="sr-only">Close</span>
        </DialogPrimitive.Close>
      </DialogPrimitive.Content>
    </DialogPortal>
  )
})
DialogContent.displayName = DialogPrimitive.Content.displayName

const DialogHeader = ({
  className,
  ...props
}: React.HTMLAttributes<HTMLDivElement>) => (
  <div
    className={cn(
      // Left-aligned at every width. The shadcn default centres below `sm`,
      // so the same dialog read as a different component on a phone.
      "flex flex-col gap-1 text-left",
      className
    )}
    {...props}
  />
)
DialogHeader.displayName = "DialogHeader"

const DialogFooter = ({
  className,
  ...props
}: React.HTMLAttributes<HTMLDivElement>) => (
  <div
    className={cn(
      "flex flex-col-reverse gap-2 sm:flex-row sm:justify-end",
      className
    )}
    {...props}
  />
)
DialogFooter.displayName = "DialogFooter"

const DialogTitle = React.forwardRef<
  React.ElementRef<typeof DialogPrimitive.Title>,
  React.ComponentPropsWithoutRef<typeof DialogPrimitive.Title>
>(({ className, ...props }, ref) => (
  <DialogPrimitive.Title
    ref={ref}
    className={cn(
      // The section-heading role: 16px/600. An overlay title is not a page
      // title, and overlay chrome does not vary by feature.
      "text-base font-semibold leading-tight tracking-tight text-fg",
      className
    )}
    {...props}
  />
))
DialogTitle.displayName = DialogPrimitive.Title.displayName

const DialogDescription = React.forwardRef<
  React.ElementRef<typeof DialogPrimitive.Description>,
  React.ComponentPropsWithoutRef<typeof DialogPrimitive.Description>
>(({ className, ...props }, ref) => (
  <DialogPrimitive.Description
    ref={ref}
    className={cn("text-sm text-muted", className)}
    {...props}
  />
))
DialogDescription.displayName = DialogPrimitive.Description.displayName

export {
  Dialog,
  DialogPortal,
  DialogOverlay,
  DialogTrigger,
  DialogClose,
  DialogContent,
  DialogHeader,
  DialogFooter,
  DialogTitle,
  DialogDescription,
}
