"use client"

import type { DndContextProps, UniqueIdentifier } from "@dnd-kit/core"
import {
  closestCenter,
  DndContext,
  DragOverlay,
  KeyboardSensor,
  PointerSensor,
  useSensor,
  useSensors,
} from "@dnd-kit/core"
import {
  horizontalListSortingStrategy,
  type NewIndexGetter,
  SortableContext,
  type SortableContextProps,
  sortableKeyboardCoordinates,
  useSortable,
  verticalListSortingStrategy,
} from "@dnd-kit/sortable"
import { CSS, type Transform } from "@dnd-kit/utilities"
import { composeEventHandlers } from "@radix-ui/primitive"
import { composeRefs } from "@radix-ui/react-compose-refs"
import * as PortalPrimitive from "@radix-ui/react-portal"
import { Primitive } from "@radix-ui/react-primitive"
import clsx from "clsx"
import type { ComponentProps, ReactNode } from "react"
import { createContext, useContext, useState } from "react"

const SortableImplContext = createContext<
  Pick<SortableProps, "getNewIndex"> &
    Required<Pick<SortableProps, "getTransformStyle">> & {
      activeId: UniqueIdentifier | null
    }
>({
  activeId: null,
  getTransformStyle: CSS.Transform.toString,
})

function useSortableImplContext() {
  const context = useContext(SortableImplContext)

  if (!context) {
    throw new Error("useSortableImplContext must be used within a <Sortable />")
  }

  return context
}

export interface SortableProps extends DndContextProps {
  onReorder?: (oldIndex: number, newIndex: number) => void
  getNewIndex?: NewIndexGetter
  getTransformStyle?: (transform: Transform | null) => React.CSSProperties["transform"]
}

export const Sortable = ({
  onDragStart,
  onDragEnd,
  onDragCancel,
  getNewIndex,
  collisionDetection = closestCenter,
  getTransformStyle = CSS.Transform.toString,
  ...props
}: SortableProps) => {
  const [activeId, setActiveId] = useState<UniqueIdentifier | null>(null)

  const sensors = useSensors(
    useSensor(PointerSensor),
    useSensor(KeyboardSensor, {
      coordinateGetter: sortableKeyboardCoordinates,
    }),
  )

  return (
    <SortableImplContext.Provider
      value={{
        activeId,
        getTransformStyle,
        getNewIndex,
      }}
    >
      <DndContext
        // @ts-expect-error Temporary fix for dnd-kit types
        onDragStart={composeEventHandlers(onDragStart, ({ active }) => setActiveId(active.id))}
        // @ts-expect-error Temporary fix for dnd-kit types
        onDragEnd={composeEventHandlers(onDragEnd, () => setActiveId(null))}
        // @ts-expect-error Temporary fix for dnd-kit types
        onDragCancel={composeEventHandlers(onDragCancel, () => setActiveId(null))}
        collisionDetection={collisionDetection}
        sensors={sensors}
        {...props}
      />
    </SortableImplContext.Provider>
  )
}

export interface SortableListProps
  extends Omit<SortableContextProps, "children">,
    ComponentProps<typeof Primitive.ul> {
  orientation?: "vertical" | "horizontal"
}

export const SortableList = ({
  orientation = "vertical",
  strategy = orientation === "vertical"
    ? verticalListSortingStrategy
    : horizontalListSortingStrategy,
  items,
  disabled,
  id,
  ...props
}: SortableListProps) => (
  <SortableContext strategy={strategy} items={items} disabled={disabled} id={id}>
    <Primitive.ul data-orientation={orientation} {...props} />
  </SortableContext>
)
SortableList.displayName = "SortableList"

export type SortableGridProps = Omit<SortableContextProps, "children"> &
  ComponentProps<typeof Primitive.div>

export const SortableGrid = ({ strategy, items, disabled, id, ...props }: SortableGridProps) => (
  <SortableContext strategy={strategy} items={items} disabled={disabled} id={id}>
    <Primitive.div {...props} />
  </SortableContext>
)
SortableGrid.displayName = "SortableGrid"

const SortableItemContext = createContext<Pick<SortableItemProps, "id" | "disabled">>({
  id: "",
  disabled: false,
})

function useSortableItemContext() {
  const context = useContext(SortableItemContext)

  if (!context) {
    throw new Error("useSortableItemContext must be used within a <SortableItem />")
  }

  return context
}

export type SortableItemProps = Omit<React.ComponentProps<typeof Primitive.div>, "id"> &
  Pick<Parameters<typeof useSortable>[0], "id" | "disabled">

export const SortableItem = ({ id, disabled, style: styleProp, ...props }: SortableItemProps) => {
  const { getTransformStyle, getNewIndex } = useSortableImplContext()

  const { attributes, setNodeRef, transform, transition, isDragging, isOver, isSorting } =
    useSortable({
      id,
      disabled,
      getNewIndex,
    })

  const style = {
    transform: getTransformStyle(transform),
    transition,
    ...styleProp,
  }

  return (
    <SortableItemContext.Provider value={{ id, disabled }}>
      <Primitive.div
        ref={composeRefs(setNodeRef as React.RefCallback<HTMLDivElement>)}
        style={style}
        data-dragging={isDragging ? true : undefined}
        data-over={isOver ? true : undefined}
        data-sorting={isSorting ? true : undefined}
        {...attributes}
        {...props}
      />
    </SortableItemContext.Provider>
  )
}
SortableItem.displayName = "SortableItem"

export const SortableItemTrigger = ({
  className,
  disabled: disabledProp,
  ...props
}: ComponentProps<typeof Primitive.button>) => {
  const { getNewIndex } = useSortableImplContext()

  const { id, disabled } = useSortableItemContext()

  const { listeners, setActivatorNodeRef, isDragging, isOver, isSorting } = useSortable({
    id,
    disabled: disabledProp ?? disabled,
    getNewIndex,
  })

  return (
    <Primitive.button
      ref={composeRefs(setActivatorNodeRef as React.RefCallback<HTMLButtonElement>)}
      data-dragging={isDragging ? true : undefined}
      data-over={isOver ? true : undefined}
      data-sorting={isSorting ? true : undefined}
      disabled={disabledProp}
      className={clsx("touch-none", className)}
      {...listeners}
      {...props}
    />
  )
}
SortableItemTrigger.displayName = "SortableItemTrigger"

export interface SortableOverlayProps extends Omit<ComponentProps<typeof DragOverlay>, "children"> {
  children?: ReactNode | ((id: UniqueIdentifier) => ReactNode)
}

export const SortableOverlay = ({ children, ...props }: SortableOverlayProps) => {
  const { activeId } = useSortableImplContext()

  return (
    <PortalPrimitive.Root>
      <DragOverlay {...props}>
        {activeId && (typeof children === "function" ? children(activeId) : children)}
      </DragOverlay>
    </PortalPrimitive.Root>
  )
}