import * as React from 'react'
import { ArrowDownIcon, ArrowUpDownIcon, ArrowUpIcon } from 'lucide-react'

import { cn } from './utils'
import { cycleSortDirection, sortRows } from '../../utils/sorting'
import { toggleSelection, toggleSelectAll } from '../../utils/selection'

/* ------------------------------------------------------------------ */
/* Types                                                               */
/* ------------------------------------------------------------------ */

export type SortDirection = 'asc' | 'desc' | null

export interface ColumnDef<TRow = Record<string, unknown>> {
  /** Unique column identifier, also used to access row data */
  key: string
  /** Column header label */
  label: string
  /** Whether the column is sortable */
  sortable?: boolean
  /** Fixed column width (CSS value) */
  width?: string
  /** Custom cell renderer */
  cell?: (value: unknown, row: TRow, index: number) => React.ReactNode
  /** Header class */
  headerClassName?: string
  /** Cell class */
  cellClassName?: string
}

export interface SortState {
  key: string
  direction: 'asc' | 'desc'
}

export interface DataTableProps<TRow = Record<string, unknown>> extends Omit<
  React.ComponentProps<'div'>,
  'children'
> {
  /** Column definitions */
  columns: ColumnDef<TRow>[]
  /** Row data */
  data: TRow[]
  /** Callback when a column is sorted */
  onSort?: (sort: SortState | null) => void
  /** Callback when a row is clicked */
  onRowClick?: (row: TRow, index: number) => void
  /** Controlled sort state */
  sort?: SortState | null
  /** Message to display when there are no rows */
  emptyMessage?: string
  /** Whether the table is loading */
  loading?: boolean
  /** Whether rows are selectable (adds checkboxes) */
  selectable?: boolean
  /** Set of selected row indices (controlled) */
  selectedRows?: Set<number>
  /** Callback when selection changes */
  onSelectionChange?: (selectedRows: Set<number>) => void
  /** Custom row className getter */
  rowClassName?: (row: TRow, index: number) => string | undefined
}

/* ------------------------------------------------------------------ */
/* Component                                                           */
/* ------------------------------------------------------------------ */

function DataTable<TRow extends Record<string, unknown> = Record<string, unknown>>({
  className,
  columns,
  data,
  onSort,
  onRowClick,
  sort: controlledSort,
  emptyMessage = 'No data available',
  loading = false,
  selectable = false,
  selectedRows: controlledSelectedRows,
  onSelectionChange,
  rowClassName,
  ...props
}: DataTableProps<TRow>) {
  // Internal sort state (uncontrolled fallback)
  const [internalSort, setInternalSort] = React.useState<SortState | null>(null)
  const sort = controlledSort !== undefined ? controlledSort : internalSort

  // Internal selection state (uncontrolled fallback)
  const [internalSelected, setInternalSelected] = React.useState<Set<number>>(new Set())
  const selectedRows = controlledSelectedRows ?? internalSelected

  const setSelectedRows = React.useCallback(
    (next: Set<number>) => {
      if (controlledSelectedRows === undefined) {
        setInternalSelected(next)
      }
      onSelectionChange?.(next)
    },
    [controlledSelectedRows, onSelectionChange]
  )

  const handleSort = React.useCallback(
    (key: string) => {
      const next = cycleSortDirection(sort?.key ?? null, sort?.direction ?? null, key)

      if (controlledSort === undefined) setInternalSort(next)
      onSort?.(next)
    },
    [sort, controlledSort, onSort]
  )

  // Sort data if using internal sort
  const sortedData = React.useMemo(() => {
    if (sort === null) return data
    const { key, direction } = sort
    return sortRows(data, direction, row => row[key as keyof TRow])
  }, [data, sort])

  const allSelected = data.length > 0 && selectedRows.size === data.length

  const toggleAll = React.useCallback(() => {
    const indices = data.map((_, i) => i)
    setSelectedRows(toggleSelectAll(indices, selectedRows))
  }, [data, selectedRows, setSelectedRows])

  const toggleRow = React.useCallback(
    (index: number) => {
      setSelectedRows(toggleSelection(index, selectedRows))
    },
    [selectedRows, setSelectedRows]
  )

  // Live region for sort announcements
  const [announcement, setAnnouncement] = React.useState('')

  const handleSortWithAnnounce = React.useCallback(
    (key: string) => {
      handleSort(key)
      const col = columns.find(c => c.key === key)
      if (col) {
        const next = cycleSortDirection(sort?.key ?? null, sort?.direction ?? null, key)
        if (next) {
          setAnnouncement(
            `Sorted by ${col.label} ${next.direction === 'asc' ? 'ascending' : 'descending'}`
          )
        } else {
          setAnnouncement(`${col.label} sort cleared`)
        }
      }
    },
    [handleSort, columns, sort]
  )

  return (
    <div data-slot="data-table" className={cn('relative w-full', className)} {...props}>
      {/* Live region for screen reader announcements */}
      <div className="sr-only" role="status" aria-live="polite" aria-atomic="true">
        {announcement}
      </div>

      {loading && (
        <div className="absolute inset-0 z-10 flex items-center justify-center rounded-xl bg-surface-primary/60 backdrop-blur-sm">
          <span className="text-foreground-muted text-sm animate-pulse">Loading...</span>
        </div>
      )}

      <div className="relative w-full overflow-x-auto rounded-xl border border-border-glass">
        <table
          role="grid"
          aria-rowcount={data.length}
          aria-colcount={columns.length}
          className="w-full caption-bottom text-sm"
        >
          <thead className="[&_tr]:border-b">
            <tr className="border-b border-border-glass">
              {selectable && (
                <th className="w-10 px-3 py-3 text-left align-middle">
                  <input
                    type="checkbox"
                    checked={allSelected}
                    onChange={toggleAll}
                    aria-label="Select all rows"
                    className="accent-primary"
                  />
                </th>
              )}
              {columns.map(col => (
                <th
                  key={col.key}
                  className={cn(
                    'text-foreground-secondary h-11 px-4 text-left align-middle text-xs font-semibold uppercase tracking-wider',
                    col.sortable === true &&
                      'cursor-pointer select-none hover:text-foreground transition-colors',
                    col.headerClassName
                  )}
                  style={col.width !== undefined ? { width: col.width } : undefined}
                  role="columnheader"
                  onClick={
                    col.sortable === true
                      ? () => {
                          handleSortWithAnnounce(col.key)
                        }
                      : undefined
                  }
                  aria-sort={
                    sort?.key === col.key
                      ? sort.direction === 'asc'
                        ? 'ascending'
                        : 'descending'
                      : undefined
                  }
                >
                  <span className="inline-flex items-center gap-1.5">
                    {col.label}
                    {col.sortable === true && (
                      <span className="text-foreground-muted inline-flex">
                        {sort?.key === col.key ? (
                          sort.direction === 'asc' ? (
                            <ArrowUpIcon className="size-3.5" />
                          ) : (
                            <ArrowDownIcon className="size-3.5" />
                          )
                        ) : (
                          <ArrowUpDownIcon className="size-3.5 opacity-40" />
                        )}
                      </span>
                    )}
                  </span>
                </th>
              ))}
            </tr>
          </thead>
          <tbody className="[&_tr:last-child]:border-0">
            {sortedData.length === 0 ? (
              <tr>
                <td
                  colSpan={columns.length + (selectable ? 1 : 0)}
                  className="px-4 py-10 text-center text-foreground-muted text-sm"
                >
                  {emptyMessage}
                </td>
              </tr>
            ) : (
              sortedData.map((row, rowIdx) => (
                <tr
                  key={rowIdx}
                  role="row"
                  aria-rowindex={rowIdx + 2}
                  className={cn(
                    'border-b border-border-glass transition-colors',
                    'hover:bg-glass-hover',
                    onRowClick !== undefined && 'cursor-pointer',
                    selectedRows.has(rowIdx) && 'bg-primary/5',
                    rowClassName?.(row, rowIdx)
                  )}
                  data-state={selectedRows.has(rowIdx) ? 'selected' : undefined}
                  onClick={
                    onRowClick !== undefined
                      ? () => {
                          onRowClick(row, rowIdx)
                        }
                      : undefined
                  }
                >
                  {selectable && (
                    <td className="w-10 px-3 py-3 align-middle">
                      <input
                        type="checkbox"
                        checked={selectedRows.has(rowIdx)}
                        onChange={e => {
                          e.stopPropagation()
                          toggleRow(rowIdx)
                        }}
                        onClick={e => {
                          e.stopPropagation()
                        }}
                        aria-label={`Select row ${String(rowIdx + 1)}`}
                        className="accent-primary"
                      />
                    </td>
                  )}
                  {columns.map(col => {
                    const cellValue = row[col.key as keyof TRow]
                    return (
                      <td
                        key={col.key}
                        role="gridcell"
                        className={cn('px-4 py-3 align-middle text-foreground', col.cellClassName)}
                      >
                        {col.cell
                          ? col.cell(cellValue, row, rowIdx)
                          : cellValue === null || cellValue === undefined
                            ? ''
                            : String(cellValue)}
                      </td>
                    )
                  })}
                </tr>
              ))
            )}
          </tbody>
        </table>
      </div>
    </div>
  )
}

export { DataTable }
