Data table

"use client";

import * as React from "react";
import { ArrowDownIcon, ArrowUpIcon, ArrowUpDownIcon } from "lucide-react";
import { Checkbox } from "@/components/ui/checkbox";
import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table";
import { cn } from "@/lib/utils";

export type DataTableColumn<T> = {
  /** Key of the row value. Also the React key. */
  key: keyof T & string;
  /** Header label. */
  header: string;
  /** Text alignment. Numbers usually go right. */
  align?: "left" | "right";
  /** Turns the header into a sort button. */
  sortable?: boolean;
  /** Value used for sorting when the raw value is not comparable. */
  sortValue?: (row: T) => string | number;
  /** Custom cell. Defaults to the raw value as text. */
  render?: (row: T) => React.ReactNode;
  /** Extra classes on the cells of this column. */
  className?: string;
};

type SortState<T> = { key: keyof T & string; direction: "asc" | "desc" } | null;

type DataTableProps<T> = Omit<React.ComponentProps<"div">, "children"> & {
  columns: DataTableColumn<T>[];
  rows: T[];
  getRowId: (row: T) => string;
  /** Shows a checkbox column and the selection count. */
  selectable?: boolean;
  /** Controlled selection, as row ids. */
  selected?: string[];
  /** Initial selection when uncontrolled. */
  defaultSelected?: string[];
  /** Called with the new list of selected ids. */
  onSelectedChange?: (ids: string[]) => void;
  /** Initial sort. */
  defaultSort?: SortState<T>;
  /** Accessible name for the table. */
  label: string;
};

function compare(a: string | number, b: string | number) {
  if (typeof a === "number" && typeof b === "number") return a - b;
  return String(a).localeCompare(String(b), undefined, { numeric: true, sensitivity: "base" });
}

/**
 * A table with sortable columns and row selection, on shadcn's table primitives.
 * Selection lives in the checkbox column, so rows stay free for links and menus.
 */
export function DataTable<T>({
  columns,
  rows,
  getRowId,
  selectable = false,
  selected,
  defaultSelected = [],
  onSelectedChange,
  defaultSort = null,
  label,
  className,
  ...props
}: DataTableProps<T>) {
  const [sort, setSort] = React.useState<SortState<T>>(defaultSort);
  const [inner, setInner] = React.useState<string[]>(defaultSelected);
  const ids = selected ?? inner;
  const chosen = React.useMemo(() => new Set(ids), [ids]);

  const sorted = React.useMemo(() => {
    if (!sort) return rows;
    const column = columns.find((c) => c.key === sort.key);
    if (!column) return rows;
    const value = (row: T) => column.sortValue?.(row) ?? (row[column.key] as unknown as string | number);
    const dir = sort.direction === "asc" ? 1 : -1;
    return [...rows].sort((a, b) => compare(value(a), value(b)) * dir);
  }, [rows, columns, sort]);

  const commit = (next: string[]) => {
    if (selected === undefined) setInner(next);
    onSelectedChange?.(next);
  };

  const rowIds = rows.map(getRowId);
  const count = rowIds.filter((id) => chosen.has(id)).length;
  const all = rows.length > 0 && count === rows.length;
  const headState: boolean | "indeterminate" = all ? true : count > 0 ? "indeterminate" : false;

  const toggleAll = () =>
    commit(all ? ids.filter((id) => !rowIds.includes(id)) : Array.from(new Set([...ids, ...rowIds])));
  const toggleRow = (id: string) => commit(chosen.has(id) ? ids.filter((x) => x !== id) : [...ids, id]);
  const toggleSort = (key: keyof T & string) =>
    setSort((cur) =>
      cur?.key !== key ? { key, direction: "asc" } : cur.direction === "asc" ? { key, direction: "desc" } : null,
    );

  return (
    <div
      data-slot="data-table"
      className={cn("w-full overflow-hidden rounded-lg border bg-background", className)}
      {...props}
    >
      <Table aria-label={label}>
        <TableHeader>
          <TableRow className="hover:bg-transparent">
            {selectable && (
              <TableHead className="h-9 w-10 pl-3">
                <Checkbox
                  aria-label="Select all rows"
                  checked={headState}
                  onCheckedChange={toggleAll}
                  className="before:hidden before:h-0.5 before:w-2 before:rounded-full before:bg-primary-foreground data-[state=indeterminate]:border-primary data-[state=indeterminate]:bg-primary data-[state=indeterminate]:before:block dark:data-[state=indeterminate]:bg-primary [&[data-state=indeterminate]_svg]:hidden"
                />
              </TableHead>
            )}
            {columns.map((column) => {
              const active = sort?.key === column.key;
              const Icon = active ? (sort.direction === "asc" ? ArrowUpIcon : ArrowDownIcon) : ArrowUpDownIcon;
              return (
                <TableHead
                  key={column.key}
                  aria-sort={active ? (sort.direction === "asc" ? "ascending" : "descending") : undefined}
                  className={cn(
                    "h-9 px-3 text-xs text-muted-foreground",
                    column.align === "right" && "text-right",
                    column.className,
                  )}
                >
                  {column.sortable ? (
                    <button
                      type="button"
                      onClick={() => toggleSort(column.key)}
                      className={cn(
                        "group/sort -mx-1.5 inline-flex h-7 items-center gap-1 rounded-sm px-1.5 transition-colors duration-100 ease-out outline-none hover:bg-muted hover:text-foreground focus-visible:ring-3 focus-visible:ring-ring/50",
                        active && "text-foreground",
                        column.align === "right" && "flex-row-reverse",
                      )}
                    >
                      {column.header}
                      <Icon
                        aria-hidden="true"
                        className={cn(
                          "size-3.5 transition-opacity duration-100 ease-out",
                          active
                            ? "opacity-100"
                            : "opacity-0 group-hover/sort:opacity-100 group-focus-visible/sort:opacity-100",
                        )}
                      />
                      <span className="sr-only">
                        {active ? (sort.direction === "asc" ? ", sorted ascending" : ", sorted descending") : ", sort"}
                      </span>
                    </button>
                  ) : (
                    column.header
                  )}
                </TableHead>
              );
            })}
          </TableRow>
        </TableHeader>
        <TableBody>
          {sorted.map((row) => {
            const id = getRowId(row);
            const isOn = chosen.has(id);
            return (
              <TableRow key={id} data-state={isOn ? "selected" : undefined}>
                {selectable && (
                  <TableCell className="h-10 w-10 py-0 pl-3">
                    <Checkbox aria-label={`Select row ${id}`} checked={isOn} onCheckedChange={() => toggleRow(id)} />
                  </TableCell>
                )}
                {columns.map((column) => (
                  <TableCell
                    key={column.key}
                    className={cn("h-10 px-3 py-0", column.align === "right" && "text-right", column.className)}
                  >
                    {column.render ? column.render(row) : String(row[column.key] ?? "")}
                  </TableCell>
                ))}
              </TableRow>
            );
          })}
        </TableBody>
      </Table>
      {selectable && (
        <p aria-live="polite" className="m-0 border-t px-3 py-2 font-mono text-xs text-muted-foreground tabular-nums">
          {count} of {rows.length} selected
        </p>
      )}
    </div>
  );
}