design-system/src/design-system/composites/DataTable/DataTable.tsx

import { useState, useMemo, Fragment } from 'react'
import styles from './DataTable.module.css'
import { Select } from '../../primitives/Select/Select'
import { Button } from '../../primitives/Button/Button'
import type { DataTableProps, Column } from './types'

type SortDir = 'asc' | 'desc'

function compareValues(a: unknown, b: unknown, dir: SortDir): number {
  const av = a == null ? '' : String(a)
  const bv = b == null ? '' : String(b)
  const cmp = av.localeCompare(bv, undefined, { numeric: true })
  return dir === 'asc' ? cmp : -cmp
}

export function DataTable<T extends { id: string }>({
  columns,
  data,
  onRowClick,
  selectedId,
  sortable = false,
  pageSize: initialPageSize = 25,
  pageSizeOptions = [10, 25, 50, 100],
  rowAccent,
  expandedContent,
  flush = false,
  fillHeight = false,
  onSortChange,
}: DataTableProps<T>) {
  const [sortKey, setSortKey] = useState<string | null>(null)
  const [sortDir, setSortDir] = useState<SortDir>('asc')
  const [page, setPage] = useState(1)
  const [pageSize, setPageSize] = useState(initialPageSize)
  const [expandedId, setExpandedId] = useState<string | null>(null)

  // When onSortChange is provided (controlled mode), skip client-side sorting
  const sorted = useMemo(() => {
    if (onSortChange) return data
    if (!sortKey) return data
    return [...data].sort((a, b) => {
      const av = (a as Record<string, unknown>)[sortKey]
      const bv = (b as Record<string, unknown>)[sortKey]
      return compareValues(av, bv, sortDir)
    })
  }, [data, sortKey, sortDir, onSortChange])

  const totalRows = sorted.length
  const totalPages = Math.max(1, Math.ceil(totalRows / pageSize))
  const safePage = Math.min(page, totalPages)
  const start = (safePage - 1) * pageSize
  const end = Math.min(start + pageSize, totalRows)
  const pageRows = sorted.slice(start, end)

  const rangeStart = totalRows === 0 ? 0 : start + 1
  const rangeInfo = `${rangeStart.toLocaleString()}–${end.toLocaleString()} of ${totalRows.toLocaleString()}`

  function handleHeaderClick(col: Column<T>) {
    if (!sortable && !col.sortable) return
    let newDir: SortDir
    if (sortKey === col.key) {
      newDir = sortDir === 'asc' ? 'desc' : 'asc'
      setSortDir(newDir)
    } else {
      newDir = 'asc'
      setSortKey(col.key)
      setSortDir(newDir)
    }
    setPage(1)
    onSortChange?.(col.key, newDir)
  }

  function handleRowClick(row: T) {
    if (expandedContent) {
      setExpandedId((prev) => (prev === row.id ? null : row.id))
    }
    onRowClick?.(row)
  }

  const pageSizeSelectOptions = pageSizeOptions.map((n) => ({
    value: String(n),
    label: String(n),
  }))

  return (
    <div className={`${styles.wrapper} ${flush ? styles.flush : ''} ${fillHeight ? styles.fillHeight : ''}`}>
      <div className={styles.scroll}>
        <table className={styles.table}>
          <thead>
            <tr>
              {columns.map((col) => {
                const isSorted = sortKey === col.key
                const canSort = sortable || col.sortable
                return (
                  <th
                    key={col.key}
                    className={[
                      styles.th,
                      isSorted ? styles.sorted : '',
                      canSort ? styles.sortable : '',
                    ]
                      .filter(Boolean)
                      .join(' ')}
                    style={col.width ? { width: col.width } : undefined}
                    onClick={() => handleHeaderClick(col)}
                    aria-sort={
                      isSorted ? (sortDir === 'asc' ? 'ascending' : 'descending') : undefined
                    }
                  >
                    {col.header}
                    {canSort && (
                      <span className={styles.sortArrow}>
                        {isSorted ? (sortDir === 'asc' ? '↑' : '↓') : '↕'}
                      </span>
                    )}
                  </th>
                )
              })}
            </tr>
          </thead>
          <tbody>
            {pageRows.map((row) => {
              const accent = rowAccent?.(row)
              const isSelected = selectedId === row.id
              const isExpanded = expandedId === row.id
              const expanded = expandedContent ? expandedContent(row) : null

              return (
                <Fragment key={row.id}>
                  <tr
                    className={[
                      styles.tr,
                      isSelected ? styles.selected : '',
                      accent === 'error' ? styles.accentError : '',
                      accent === 'warning' ? styles.accentWarning : '',
                    ]
                      .filter(Boolean)
                      .join(' ')}
                    onClick={() => handleRowClick(row)}
                  >
                    {columns.map((col) => {
                      const raw = (row as Record<string, unknown>)[col.key]
                      return (
                        <td key={col.key} className={styles.td}>
                          {col.render
                            ? col.render(raw, row)
                            : (raw as string | number | boolean | null | undefined)}
                        </td>
                      )
                    })}
                  </tr>
                  {isExpanded && expanded != null && (
                    <tr className={styles.expandedRow}>
                      <td colSpan={columns.length} className={styles.expandedCell}>
                        {expanded}
                      </td>
                    </tr>
                  )}
                </Fragment>
              )
            })}
            {pageRows.length === 0 && (
              <tr>
                <td colSpan={columns.length} className={styles.empty}>
                  No data
                </td>
              </tr>
            )}
          </tbody>
        </table>
      </div>

      <div className={styles.footer}>
        <span className={styles.rangeInfo}>{rangeInfo}</span>
        <div className={styles.paginationRight}>
          <span className={styles.pageSizeLabel}>Rows:</span>
          <Select
            options={pageSizeSelectOptions}
            value={String(pageSize)}
            onChange={(e) => {
              setPageSize(Number(e.target.value))
              setPage(1)
            }}
            className={styles.pageSizeSelect}
          />
          <div className={styles.pagination}>
            <Button
              size="sm"
              variant="secondary"
              disabled={safePage <= 1}
              onClick={() => setPage((p) => p - 1)}
              aria-label="Previous page"
            >
              ‹
            </Button>
            <span className={styles.pageNum}>
              {safePage} / {totalPages}
            </span>
            <Button
              size="sm"
              variant="secondary"
              disabled={safePage >= totalPages}
              onClick={() => setPage((p) => p + 1)}
              aria-label="Next page"
            >
              ›
            </Button>
          </div>
        </div>
      </div>
    </div>
  )
}