Table

"use client"

import * as React from "react"
import {
  type ColumnDef,
  type ColumnSizingState,
  type RowSelectionState,
  type SortingState,
  flexRender,
  getCoreRowModel,
  getSortedRowModel,
  useReactTable,
} from "@tanstack/react-table"
import { useVirtualizer, useWindowVirtualizer } from "@tanstack/react-virtual"
import { ChevronDown } from "lucide-react"

import { Checkbox } from "@/components/ui/checkbox"
import { cn } from "@/lib/utils"

type Align = "left" | "center" | "right"
type SortDirection = "asc" | "desc"
type ColumnType =
  | "text"
  | "longText"
  | "shortText"
  | "number"
  | "currency"
  | "date"
  | "status"
  | "actions"

export type EmsteinUITableColumn<TData> = {
  key: Extract<keyof TData, string> | string
  label: string
  type?: ColumnType
  grow?: boolean | number
  width?: number
  minWidth?: number
  sortable?: boolean
  resizable?: boolean
  align?: Align
  render?: (row: TData) => React.ReactNode
  footer?:
    | React.ReactNode
    | ((context: { data: TData[]; rows: TData[] }) => React.ReactNode)
  className?: string
  footerClassName?: string
}

export type EmsteinUITableSelection = {
  selectedRows?: string[]
  onSelectedRowsChange?: (rows: string[]) => void
  defaultSelectedRows?: string[]
}

export type EmsteinUITableProps<TData> = {
  data: TData[]
  columns: EmsteinUITableColumn<TData>[]
  getRowId?: (row: TData, index: number) => string
  selection?: EmsteinUITableSelection
  resizable?: boolean
  sortable?: boolean
  showUnsortedSortIcon?: boolean
  initialSort?: { key: string; direction?: SortDirection } | false
  scrollMode?: "page" | "table"
  tableHeight?: number
  rowHeight?: number
  emptyText?: string
  caption?: React.ReactNode
  className?: string
  containerClassName?: string
  tableClassName?: string
  headerClassName?: string
  bodyClassName?: string
  footerClassName?: string
  captionClassName?: string
}

export type TablePrimitiveProps = React.ComponentProps<"table"> & {
  containerClassName?: string
}

const alignClass: Record<Align, string> = {
  left: "justify-start text-left",
  center: "justify-center text-center",
  right: "justify-end text-right",
}

const columnTypeSizes: Record<
  ColumnType,
  { width: number; minWidth: number; grow: number }
> = {
  text: { width: 180, minWidth: 120, grow: 1 },
  longText: { width: 280, minWidth: 180, grow: 2 },
  shortText: { width: 140, minWidth: 96, grow: 0.5 },
  number: { width: 104, minWidth: 80, grow: 0 },
  currency: { width: 124, minWidth: 96, grow: 0 },
  date: { width: 132, minWidth: 112, grow: 0 },
  status: { width: 128, minWidth: 104, grow: 0 },
  actions: { width: 72, minWidth: 56, grow: 0 },
}

function getColumnSizes<TData>(column: EmsteinUITableColumn<TData>) {
  const preset = columnTypeSizes[column.type ?? "text"]

  return {
    width: column.width ?? preset.width,
    minWidth: column.minWidth ?? preset.minWidth,
  }
}

function getColumnGrow<TData>(column: EmsteinUITableColumn<TData>) {
  if (typeof column.grow === "number") {
    return Math.max(0, column.grow)
  }

  if (column.grow === true) {
    return 1
  }

  if (column.grow === false) {
    return 0
  }

  return columnTypeSizes[column.type ?? "text"].grow
}

function getValue<TData>(row: TData, key: string) {
  return (row as Record<string, unknown>)[key]
}

function toRowSelection(ids: string[]): RowSelectionState {
  return ids.reduce<RowSelectionState>((state, id) => {
    state[id] = true
    return state
  }, {})
}

function getNextValue<TValue>(
  updater: TValue | ((old: TValue) => TValue),
  old: TValue
) {
  return typeof updater === "function"
    ? (updater as (value: TValue) => TValue)(old)
    : updater
}

export function TableHeader({
  className,
  ...props
}: React.ComponentProps<"thead">) {
  return (
    <thead
      data-slot="table-header"
      className={cn("[&_tr]:border-b", className)}
      {...props}
    />
  )
}

export function TableBody({
  className,
  ...props
}: React.ComponentProps<"tbody">) {
  return (
    <tbody
      data-slot="table-body"
      className={cn("[&_tr:last-child]:border-0", className)}
      {...props}
    />
  )
}

export function TableFooter({
  className,
  ...props
}: React.ComponentProps<"tfoot">) {
  return (
    <tfoot
      data-slot="table-footer"
      className={cn(
        "border-t bg-muted/50 font-medium [&>tr]:last:border-b-0",
        className
      )}
      {...props}
    />
  )
}

export function TableRow({ className, ...props }: React.ComponentProps<"tr">) {
  return (
    <tr
      data-slot="table-row"
      className={cn(
        "border-b transition-colors hover:bg-muted/50 has-aria-expanded:bg-muted/50 data-[state=selected]:bg-muted",
        className
      )}
      {...props}
    />
  )
}

export function TableHead({ className, ...props }: React.ComponentProps<"th">) {
  return (
    <th
      data-slot="table-head"
      className={cn(
        "h-12 px-3 text-left align-middle text-xs font-medium uppercase tracking-wider whitespace-nowrap text-muted-foreground [&:has([role=checkbox])]:pr-0",
        className
      )}
      {...props}
    />
  )
}

export function TableCell({ className, ...props }: React.ComponentProps<"td">) {
  return (
    <td
      data-slot="table-cell"
      className={cn(
        "p-3 align-middle whitespace-nowrap [&:has([role=checkbox])]:pr-0",
        className
      )}
      {...props}
    />
  )
}

export function TableCaption({
  className,
  ...props
}: React.ComponentProps<"caption">) {
  return (
    <caption
      data-slot="table-caption"
      className={cn("mt-4 text-sm text-muted-foreground", className)}
      {...props}
    />
  )
}

function PrimitiveTable({
  className,
  containerClassName,
  ...props
}: TablePrimitiveProps) {
  return (
    <div
      data-slot="table-container"
      className={cn("relative w-full overflow-x-auto", containerClassName)}
    >
      <table
        data-slot="table"
        className={cn("w-full caption-bottom text-sm", className)}
        {...props}
      />
    </div>
  )
}

function DataTable<TData>({
  data,
  columns,
  getRowId,
  selection,
  resizable = true,
  sortable = true,
  showUnsortedSortIcon = false,
  initialSort,
  scrollMode = "page",
  tableHeight = 560,
  rowHeight = 56,
  emptyText = "No results.",
  caption,
  className,
  containerClassName,
  tableClassName,
  headerClassName,
  bodyClassName,
  footerClassName,
  captionClassName,
}: EmsteinUITableProps<TData>) {
  const isSelectionControlled = selection?.selectedRows !== undefined
  const [internalSelectedRows, setInternalSelectedRows] = React.useState(
    selection?.defaultSelectedRows ?? []
  )
  const selectedRows = selection?.selectedRows ?? internalSelectedRows
  const rowSelection = React.useMemo(
    () => toRowSelection(selectedRows),
    [selectedRows]
  )
  const hasSelection = selection !== undefined

  const initialSorting = React.useMemo<SortingState>(() => {
    if (!sortable || initialSort === false) {
      return []
    }

    if (initialSort?.key) {
      return [{ id: initialSort.key, desc: initialSort.direction === "desc" }]
    }

    const firstSortableColumn = columns.find((column) => column.sortable !== false)

    return firstSortableColumn
      ? [{ id: String(firstSortableColumn.key), desc: false }]
      : []
  }, [columns, initialSort, sortable])

  const [sorting, setSorting] = React.useState<SortingState>(initialSorting)
  const [columnSizing, setColumnSizing] = React.useState<ColumnSizingState>({})

  const tableColumns = React.useMemo<ColumnDef<TData>[]>(
    () => [
      ...(hasSelection
        ? [
            {
              id: "_select",
              size: 44,
              minSize: 44,
              maxSize: 44,
              enableSorting: false,
              enableResizing: false,
              header: ({ table }) => (
                <Checkbox
                  aria-label="Select all rows"
                  checked={
                    table.getIsAllRowsSelected()
                      ? true
                      : table.getIsSomeRowsSelected()
                        ? "indeterminate"
                        : false
                  }
                  onCheckedChange={(value) =>
                    table.toggleAllRowsSelected(Boolean(value))
                  }
                />
              ),
              cell: ({ row }) => (
                <Checkbox
                  aria-label="Select row"
                  checked={row.getIsSelected()}
                  onCheckedChange={(value) => row.toggleSelected(Boolean(value))}
                />
              ),
            } satisfies ColumnDef<TData>,
          ]
        : []),
      ...columns.map(
        (column): ColumnDef<TData> => {
          const sizes = getColumnSizes(column)

          return {
            id: String(column.key),
            accessorFn: (row) => getValue(row, String(column.key)),
            size: sizes.width,
            minSize: sizes.minWidth,
            enableSorting: sortable && column.sortable !== false,
            enableResizing: resizable && column.resizable !== false,
            header: column.label,
            cell: ({ row }) => {
              const value = getValue(row.original, String(column.key))

              return column.render
                ? column.render(row.original)
                : String(value ?? "")
            },
            meta: {
              align: column.align ?? "left",
              className: column.className,
              grow: getColumnGrow(column),
            },
          }
        }
      ),
    ],
    [columns, hasSelection, resizable, sortable]
  )

  const table = useReactTable({
    data,
    columns: tableColumns,
    state: {
      sorting,
      columnSizing,
      rowSelection,
    },
    defaultColumn: {
      minSize: 96,
      size: 160,
    },
    getRowId: (row, index) => getRowId?.(row, index) ?? String(index),
    enableRowSelection: hasSelection,
    enableSortingRemoval: false,
    columnResizeMode: "onChange",
    onSortingChange: setSorting,
    onColumnSizingChange: setColumnSizing,
    onRowSelectionChange: (updater) => {
      const next = getNextValue(updater, rowSelection)
      const ids = Object.keys(next).filter((id) => next[id])

      if (!isSelectionControlled) {
        setInternalSelectedRows(ids)
      }

      selection?.onSelectedRowsChange?.(ids)
    },
    getCoreRowModel: getCoreRowModel(),
    getSortedRowModel: getSortedRowModel(),
  })

  const rows = table.getRowModel().rows
  const renderedRows = React.useMemo(
    () => rows.map((row) => row.original),
    [rows]
  )
  const tableRootRef = React.useRef<HTMLDivElement>(null)
  const tableScrollRef = React.useRef<HTMLDivElement>(null)
  const [scrollMargin, setScrollMargin] = React.useState(0)
  const [containerWidth, setContainerWidth] = React.useState(0)

  React.useLayoutEffect(() => {
    if (scrollMode !== "page" || !tableRootRef.current) {
      return
    }

    const measure = () =>
      setScrollMargin(
        tableRootRef.current
          ? tableRootRef.current.getBoundingClientRect().top + window.scrollY
          : 0
      )

    measure()
    window.addEventListener("resize", measure)

    return () => window.removeEventListener("resize", measure)
  }, [scrollMode])

  React.useLayoutEffect(() => {
    if (!tableRootRef.current) {
      return
    }

    const root = tableRootRef.current
    const updateWidth = () => setContainerWidth(root.clientWidth)
    const observer = new ResizeObserver(updateWidth)

    updateWidth()
    observer.observe(root)

    return () => observer.disconnect()
  }, [])

  const windowVirtualizer = useWindowVirtualizer({
    count: rows.length,
    estimateSize: () => rowHeight,
    overscan: 12,
    scrollMargin,
    enabled: scrollMode === "page",
    initialRect: { width: 1200, height: 800 },
  })

  const tableVirtualizer = useVirtualizer({
    count: rows.length,
    estimateSize: () => rowHeight,
    overscan: 12,
    getScrollElement: () => tableScrollRef.current,
    enabled: scrollMode === "table",
  })

  const virtualizer =
    scrollMode === "page" ? windowVirtualizer : tableVirtualizer
  const virtualRows = virtualizer.getVirtualItems()
  const totalSize = virtualizer.getTotalSize()
  const leafHeaders = table.getFlatHeaders()
  const visibleLeafHeaders = table.getVisibleLeafColumns()
  const rawTableWidth = table.getTotalSize()
  const tableWidth = Math.max(rawTableWidth, containerWidth)
  const extraWidth = Math.max(0, tableWidth - rawTableWidth)
  const growByColumn = React.useMemo(() => {
    return Object.fromEntries(
      leafHeaders.map((header) => {
        const meta = header.column.columnDef.meta as
          | { grow?: number }
          | undefined

        return [header.column.id, meta?.grow ?? 0]
      })
    )
  }, [leafHeaders])
  const totalGrow = React.useMemo(
    () => Object.values(growByColumn).reduce((total, grow) => total + grow, 0),
    [growByColumn]
  )
  const getColumnWidth = React.useCallback(
    (columnId: string, size: number) => {
      if (!extraWidth || !totalGrow) {
        return size
      }

      return size + extraWidth * ((growByColumn[columnId] ?? 0) / totalGrow)
    },
    [extraWidth, growByColumn, totalGrow]
  )
  const hasFooter = columns.some((column) => column.footer !== undefined)

  function resizeWithNeighbor(
    headerIndex: number,
    event: React.PointerEvent<HTMLButtonElement>
  ) {
    const header = leafHeaders[headerIndex]
    const nextHeader = leafHeaders[headerIndex + 1]

    if (!header?.column.getCanResize() || !nextHeader) {
      return
    }

    event.preventDefault()
    event.currentTarget.setPointerCapture(event.pointerId)

    const startX = event.clientX
    const startSize = getColumnWidth(header.column.id, header.column.getSize())
    const nextStartSize = getColumnWidth(
      nextHeader.column.id,
      nextHeader.column.getSize()
    )
    const minSize = header.column.columnDef.minSize ?? 96
    const nextMinSize = nextHeader.column.columnDef.minSize ?? 96

    const onPointerMove = (moveEvent: PointerEvent) => {
      const delta = moveEvent.clientX - startX
      const clampedDelta = Math.min(
        Math.max(delta, minSize - startSize),
        nextStartSize - nextMinSize
      )

      setColumnSizing((current) => ({
        ...current,
        [header.column.id]: startSize + clampedDelta,
        [nextHeader.column.id]: nextStartSize - clampedDelta,
      }))
    }

    const onPointerUp = () => {
      window.removeEventListener("pointermove", onPointerMove)
      window.removeEventListener("pointerup", onPointerUp)
    }

    window.addEventListener("pointermove", onPointerMove)
    window.addEventListener("pointerup", onPointerUp)
  }

  return (
    <div
      ref={tableRootRef}
      className={cn(
        "relative w-full border",
        scrollMode === "page" && "overflow-x-auto overflow-y-clip",
        scrollMode === "table" && "overflow-hidden",
        className,
        containerClassName
      )}
    >
      <div
        ref={tableScrollRef}
        className={cn(
          scrollMode === "table" && "overflow-auto",
          scrollMode === "page" && "overflow-visible"
        )}
        style={scrollMode === "table" ? { height: tableHeight } : undefined}
      >
        <table
          className={cn("grid w-full caption-bottom text-sm", tableClassName)}
          style={{ minWidth: rawTableWidth, width: tableWidth }}
        >
          {caption ? (
            <TableCaption className={captionClassName}>{caption}</TableCaption>
          ) : null}
          <TableHeader className={cn("grid bg-muted", headerClassName)}>
            {table.getHeaderGroups().map((headerGroup) => (
              <TableRow key={headerGroup.id} className="flex">
                {headerGroup.headers.map((header, headerIndex) => {
                  const canSort = header.column.getCanSort()
                  const sortState = header.column.getIsSorted()
                  const meta = header.column.columnDef.meta as
                    | { align?: Align }
                    | undefined
                  const align = meta?.align ?? "left"
                  const showSortIcon =
                    sortState || (showUnsortedSortIcon && canSort)
                  const renderedHeader = flexRender(
                    header.column.columnDef.header,
                    header.getContext()
                  )

                  return (
                    <TableHead
                      key={header.id}
                      className={cn(
                        "group relative flex items-center px-3",
                        canSort && "cursor-pointer",
                        alignClass[align]
                      )}
                      style={{
                        width: getColumnWidth(header.column.id, header.getSize()),
                        minWidth: header.column.columnDef.minSize,
                      }}
                    >
                      {canSort ? (
                        <button
                          type="button"
                          aria-label={`Sort ${
                            typeof header.column.columnDef.header === "string"
                              ? header.column.columnDef.header
                              : "column"
                          }`}
                          className="absolute inset-0 z-0 cursor-pointer outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
                          onClick={header.column.getToggleSortingHandler()}
                        />
                      ) : null}
                      {headerIndex > 0 ? (
                        <span
                          aria-hidden="true"
                          className="pointer-events-none absolute inset-y-2 left-0 z-10 w-px bg-ring opacity-0 transition-opacity group-hover:opacity-100"
                        />
                      ) : null}
                      {headerIndex < leafHeaders.length - 1 ? (
                        <span
                          aria-hidden="true"
                          className="pointer-events-none absolute inset-y-2 right-0 z-10 w-px bg-ring opacity-0 transition-opacity group-hover:opacity-100"
                        />
                      ) : null}
                      <div className="pointer-events-none relative z-10 inline-flex min-w-0 items-center gap-1.5 text-xs font-medium uppercase tracking-wider">
                        <span className="truncate">
                          {renderedHeader}
                        </span>
                        {showSortIcon ? (
                          <ChevronDown
                            data-icon="inline-end"
                            className={cn(
                              "transition-transform duration-200 ease-out",
                              sortState === "asc" && "rotate-180",
                              !sortState && "opacity-60"
                            )}
                          />
                        ) : null}
                      </div>
                      {resizable &&
                      header.column.getCanResize() &&
                      headerIndex < leafHeaders.length - 1 ? (
                        <button
                          type="button"
                          aria-label="Resize column"
                          className="absolute inset-y-0 right-0 z-20 w-3 cursor-col-resize outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-inset"
                          onPointerDown={(event) =>
                            resizeWithNeighbor(headerIndex, event)
                          }
                        />
                      ) : null}
                    </TableHead>
                  )
                })}
              </TableRow>
            ))}
          </TableHeader>
          <TableBody
            className={cn("relative grid", bodyClassName)}
            style={{ height: totalSize }}
          >
            {virtualRows.length ? (
              virtualRows.map((virtualRow) => {
                const row = rows[virtualRow.index]

                return (
                  <TableRow
                    key={row.id}
                    data-state={row.getIsSelected() && "selected"}
                    className="absolute flex w-full"
                    style={{
                      height: virtualRow.size,
                      transform: `translateY(${
                        scrollMode === "page"
                          ? virtualRow.start - scrollMargin
                          : virtualRow.start
                      }px)`,
                    }}
                  >
                    {row.getVisibleCells().map((cell) => {
                      const meta = cell.column.columnDef.meta as
                        | { align?: Align; className?: string }
                        | undefined
                      const align = meta?.align ?? "left"

                      return (
                        <TableCell
                          key={cell.id}
                          className={cn(
                            "flex items-center px-3",
                            alignClass[align],
                            meta?.className
                          )}
                          style={{
                            width: getColumnWidth(
                              cell.column.id,
                              cell.column.getSize()
                            ),
                            minWidth: cell.column.columnDef.minSize,
                          }}
                        >
                          <span className="truncate">
                            {flexRender(
                              cell.column.columnDef.cell,
                              cell.getContext()
                            )}
                          </span>
                        </TableCell>
                      )
                    })}
                  </TableRow>
                )
              })
            ) : (
              <TableRow className="flex">
                <TableCell className="flex h-24 w-full items-center justify-center text-muted-foreground">
                  {emptyText}
                </TableCell>
              </TableRow>
            )}
          </TableBody>
          {hasFooter ? (
            <TableFooter className={cn("grid bg-muted", footerClassName)}>
              <TableRow className="flex">
                {hasSelection ? (
                  <TableCell
                    className="flex items-center px-3"
                    style={{
                      width: getColumnWidth(
                        visibleLeafHeaders[0].id,
                        visibleLeafHeaders[0].getSize()
                      ),
                      minWidth: visibleLeafHeaders[0].columnDef.minSize,
                    }}
                  />
                ) : null}
                {columns.map((column) => {
                  const tableColumn = table.getColumn(String(column.key))

                  if (!tableColumn) {
                    return null
                  }

                  const align = column.align ?? "left"
                  const footer =
                    typeof column.footer === "function"
                      ? column.footer({ data, rows: renderedRows })
                      : column.footer

                  return (
                    <TableCell
                      key={String(column.key)}
                      className={cn(
                        "flex items-center px-3",
                        alignClass[align],
                        column.footerClassName
                      )}
                      style={{
                        width: getColumnWidth(
                          tableColumn.id,
                          tableColumn.getSize()
                        ),
                        minWidth: tableColumn.columnDef.minSize,
                      }}
                    >
                      <span className="truncate">{footer}</span>
                    </TableCell>
                  )
                })}
              </TableRow>
            </TableFooter>
          ) : null}
        </table>
      </div>
    </div>
  )
}

export function Table<TData>(
  props: EmsteinUITableProps<TData> | TablePrimitiveProps
) {
  if ("data" in props && "columns" in props) {
    return <DataTable {...props} />
  }

  return <PrimitiveTable {...props} />
}