import * as React from 'react'
import { Button } from './Button'
import { cn } from './cva'
import { ErrorState } from './ErrorState'
import { Input } from './Input'
import { Pagination } from './Pagination'
import { Skeleton } from './Skeleton'
import { Table } from './Table'

export interface DataTableColumn<T> {
  id: string
  header: React.ReactNode
  cell?(row: T): React.ReactNode
  sortable?: boolean
  sortValue?(row: T): string | number | null
  searchValue?(row: T): string
  resizable?: boolean
  width?: number
  minWidth?: number
}

export interface DataTableSort {
  columnId: string
  direction: 'asc' | 'desc'
}

export interface DataTableSortingOptions {
  sort?: DataTableSort | null
  defaultSort?: DataTableSort | null
  onSortChange?(sort: DataTableSort | null): void
  manual?: boolean
}

export interface DataTableSearchOptions {
  query?: string
  defaultQuery?: string
  onQueryChange?(query: string): void
  manual?: boolean
  label?: string
}

export interface DataTableColumnResizingOptions {
  widths?: Record<string, number>
  defaultWidths?: Record<string, number>
  onWidthsChange?(widths: Record<string, number>): void
}

interface DataTablePaginationBaseOptions {
  page?: number
  defaultPage?: number
  onPageChange?(page: number): void
  label?: string
}

export type DataTablePaginationOptions = DataTablePaginationBaseOptions &
  (
    | { manual: true; pageCount: number; pageSize?: never }
    | { manual?: false; pageSize?: number; pageCount?: never }
  )

export interface DataTableCapability {
  readonly kind: 'sorting' | 'search' | 'resizing' | 'pagination'
}

interface CapabilityImpl extends DataTableCapability {
  readonly options:
    | DataTableSortingOptions
    | DataTableSearchOptions
    | DataTableColumnResizingOptions
    | DataTablePaginationOptions
}

export function withSorting(options: DataTableSortingOptions = {}): DataTableCapability {
  const capability: CapabilityImpl = { kind: 'sorting', options }
  return capability
}

export function withSearch(options: DataTableSearchOptions = {}): DataTableCapability {
  const capability: CapabilityImpl = { kind: 'search', options }
  return capability
}

export function withColumnResizing(options: DataTableColumnResizingOptions = {}): DataTableCapability {
  const capability: CapabilityImpl = { kind: 'resizing', options }
  return capability
}

export function withPagination(options: DataTablePaginationOptions): DataTableCapability {
  const capability: CapabilityImpl = { kind: 'pagination', options }
  return capability
}

export interface DataTableProps<T> {
  caption: string
  columns: Array<DataTableColumn<T>>
  rows: readonly T[]
  getRowId(row: T): string
  capabilities?: DataTableCapability[]
  loading?: boolean
  error?: string | null
  onRetry?(): void
  emptyState?: React.ReactNode
  className?: string
}

const MIN_WIDTH_DEFAULT = 48
const KEYBOARD_RESIZE_STEP = 16

function useControllableState<S>(
  controlled: S | undefined,
  defaultValue: S,
  onChange?: (next: S) => void,
): [S, (next: S) => void] {
  const [internal, setInternal] = React.useState<S>(defaultValue)
  const value = controlled !== undefined ? controlled : internal
  const set = React.useCallback(
    (next: S) => {
      if (controlled === undefined) setInternal(next)
      onChange?.(next)
    },
    [controlled, onChange],
  )
  return [value, set]
}

function pickCapability<O>(capabilities: DataTableCapability[] | undefined, kind: DataTableCapability['kind']): O | undefined {
  if (!capabilities) return undefined
  for (let i = capabilities.length - 1; i >= 0; i--) {
    const capability = capabilities[i]
    if (capability?.kind === kind && 'options' in capability) return (capability as CapabilityImpl).options as O
  }
  return undefined
}

function accessorValue<T>(row: T, column: DataTableColumn<T>): unknown {
  return (row as Record<string, unknown>)[column.id]
}

function defaultCell<T>(row: T, column: DataTableColumn<T>): React.ReactNode {
  const value = accessorValue(row, column)
  return value === null || value === undefined ? '' : String(value)
}

function sortKey<T>(row: T, column: DataTableColumn<T>): string | number | null {
  if (column.sortValue) return column.sortValue(row)
  const value = accessorValue(row, column)
  if (value === null || value === undefined) return null
  return typeof value === 'number' ? value : String(value)
}

function searchKey<T>(row: T, column: DataTableColumn<T>): string {
  if (column.searchValue) return column.searchValue(row)
  const value = accessorValue(row, column)
  return value === null || value === undefined ? '' : String(value)
}

function compareKeys(a: string | number | null, b: string | number | null): number {
  if (typeof a === 'number' && typeof b === 'number') return a - b
  return String(a).localeCompare(String(b))
}

interface ResizeHandleProps {
  columnLabel: string
  width: number
  minWidth: number
  onResize(width: number): void
}

function ResizeHandle({ columnLabel, width, minWidth, onResize }: ResizeHandleProps) {
  const dragOrigin = React.useRef<{ pointerX: number; startWidth: number } | null>(null)

  return (
    <span
      role="separator"
      aria-orientation="vertical"
      aria-label={`Resize ${columnLabel} column`}
      aria-valuenow={Math.round(width)}
      aria-valuemin={minWidth}
      tabIndex={0}
      className="absolute inset-y-0 right-0 w-2 cursor-col-resize touch-none select-none focus-visible:outline focus-visible:outline-accent"
      onKeyDown={(event) => {
        if (event.key === 'ArrowLeft' || event.key === 'ArrowRight') {
          event.preventDefault()
          const delta = event.key === 'ArrowRight' ? KEYBOARD_RESIZE_STEP : -KEYBOARD_RESIZE_STEP
          onResize(Math.max(minWidth, width + delta))
        }
      }}
      onPointerDown={(event) => {
        event.preventDefault()
        dragOrigin.current = { pointerX: event.clientX, startWidth: width }
        event.currentTarget.setPointerCapture(event.pointerId)
      }}
      onPointerMove={(event) => {
        if (!dragOrigin.current) return
        const next = dragOrigin.current.startWidth + (event.clientX - dragOrigin.current.pointerX)
        onResize(Math.max(minWidth, next))
      }}
      onPointerUp={(event) => {
        dragOrigin.current = null
        event.currentTarget.releasePointerCapture(event.pointerId)
      }}
    />
  )
}

export function DataTable<T>({
  caption,
  columns,
  rows,
  getRowId,
  capabilities,
  loading = false,
  error = null,
  onRetry,
  emptyState = 'No results.',
  className,
}: DataTableProps<T>) {
  const searchId = React.useId()
  const sorting = pickCapability<DataTableSortingOptions>(capabilities, 'sorting')
  const search = pickCapability<DataTableSearchOptions>(capabilities, 'search')
  const resizing = pickCapability<DataTableColumnResizingOptions>(capabilities, 'resizing')
  const pagination = pickCapability<DataTablePaginationOptions>(capabilities, 'pagination')

  const [sort, setSort] = useControllableState<DataTableSort | null>(
    sorting?.sort,
    sorting?.defaultSort ?? null,
    sorting?.onSortChange,
  )
  const [query, setQuery] = useControllableState<string>(
    search?.query,
    search?.defaultQuery ?? '',
    search?.onQueryChange,
  )
  const [page, setPage] = useControllableState<number>(
    pagination?.page,
    pagination?.defaultPage ?? 1,
    pagination?.onPageChange,
  )
  const [widths, setWidths] = useControllableState<Record<string, number>>(
    resizing?.widths,
    resizing?.defaultWidths ?? {},
    resizing?.onWidthsChange,
  )

  const filtered = React.useMemo(() => {
    if (!search || search.manual || query.trim() === '') return rows
    const needle = query.trim().toLowerCase()
    return rows.filter((row) => columns.some((column) => searchKey(row, column).toLowerCase().includes(needle)))
  }, [rows, columns, query, search])

  const sorted = React.useMemo(() => {
    if (!sorting || sorting.manual || !sort) return filtered
    const column = columns.find((c) => c.id === sort.columnId)
    if (!column) return filtered
    const direction = sort.direction === 'asc' ? 1 : -1
    return filtered
      .map((row, index) => ({ row, index }))
      .sort((a, b) => {
        const keyA = sortKey(a.row, column)
        const keyB = sortKey(b.row, column)
        if (keyA === null || keyB === null) {
          if (keyA === null && keyB === null) return a.index - b.index
          return keyA === null ? 1 : -1
        }
        const order = compareKeys(keyA, keyB)
        return order !== 0 ? order * direction : a.index - b.index
      })
      .map((entry) => entry.row)
  }, [filtered, columns, sort, sorting])

  const clientPaginated = pagination !== undefined && !pagination.manual && pagination.pageSize !== undefined
  const pageCount = pagination?.manual
    ? pagination.pageCount
    : clientPaginated && pagination?.pageSize !== undefined
      ? Math.max(1, Math.ceil(sorted.length / pagination.pageSize))
      : 1
  const effectivePage = clientPaginated ? Math.min(page, pageCount) : page
  const paginated = React.useMemo(() => {
    if (!clientPaginated || pagination?.pageSize === undefined) return sorted
    return sorted.slice((effectivePage - 1) * pagination.pageSize, effectivePage * pagination.pageSize)
  }, [sorted, effectivePage, clientPaginated, pagination?.pageSize])

  const showPagination = pagination !== undefined && (pagination.manual === true || pageCount > 1)

  function resetToFirstPage() {
    if (clientPaginated) setPage(1)
  }

  function cycleSort(columnId: string) {
    if (!sort || sort.columnId !== columnId) {
      setSort({ columnId, direction: 'asc' })
    } else {
      setSort(sort.direction === 'asc' ? { columnId, direction: 'desc' } : null)
    }
    resetToFirstPage()
  }

  function updateQuery(next: string) {
    setQuery(next)
    resetToFirstPage()
  }

  const anyWidth =
    resizing !== undefined && columns.some((column) => widths[column.id] !== undefined || column.width !== undefined)

  function columnLabel(column: DataTableColumn<T>): string {
    return typeof column.header === 'string' ? column.header : column.id
  }

  return (
    <div className={cn('flex flex-col gap-3', className)} aria-busy={loading || undefined}>
      {search ? (
        <div className="flex flex-col gap-1">
          <label htmlFor={searchId} className="text-sm font-semibold text-fg">
            {search.label ?? 'Search'}
          </label>
          <Input
            id={searchId}
            type="search"
            size="sm"
            value={query}
            onChange={(event) => updateQuery(event.currentTarget.value)}
          />
        </div>
      ) : null}

      {search ? (
        <p role="status" aria-live="polite" className="sr-only">
          {loading || error ? '' : `${sorted.length} results`}
        </p>
      ) : null}

      {error ? (
        <ErrorState
          variant="inline"
          title="Data couldn't load"
          message={error}
          action={
            onRetry ? (
              <Button tone="surface" radius="sm" size="sm" type="button" onClick={onRetry}>
                Try again
              </Button>
            ) : undefined
          }
        />
      ) : loading ? (
        <div className="flex flex-col gap-2" aria-hidden="true">
          <Skeleton variant="block" />
          <Skeleton variant="block" />
        </div>
      ) : paginated.length === 0 ? (
        <div className="text-sm text-fg-muted">{emptyState}</div>
      ) : (
        <div className="overflow-x-auto">
          <Table caption={caption} className={anyWidth ? 'table-fixed' : undefined}>
            {anyWidth ? (
              <colgroup>
                {columns.map((column) => {
                  const width = widths[column.id] ?? column.width
                  return <col key={column.id} style={width !== undefined ? { width: `${width}px` } : undefined} />
                })}
              </colgroup>
            ) : null}
            <Table.Head>
              <Table.Row>
                {columns.map((column) => {
                  const sortEnabled = sorting !== undefined && column.sortable === true
                  const isSorted = sort?.columnId === column.id
                  const resizeEnabled = resizing !== undefined && column.resizable === true
                  const width = widths[column.id] ?? column.width ?? 0
                  return (
                    <Table.Th
                      key={column.id}
                      scope="col"
                      className={resizeEnabled ? 'relative' : undefined}
                      sort={sortEnabled ? (isSorted && sort ? sort.direction : 'none') : undefined}
                      onSort={sortEnabled ? () => cycleSort(column.id) : undefined}
                    >
                      {column.header}
                      {resizeEnabled ? (
                        <ResizeHandle
                          columnLabel={columnLabel(column)}
                          width={width || (column.minWidth ?? MIN_WIDTH_DEFAULT)}
                          minWidth={column.minWidth ?? MIN_WIDTH_DEFAULT}
                          onResize={(next) => setWidths({ ...widths, [column.id]: next })}
                        />
                      ) : null}
                    </Table.Th>
                  )
                })}
              </Table.Row>
            </Table.Head>
            <Table.Body>
              {paginated.map((row) => (
                <Table.Row key={getRowId(row)}>
                  {columns.map((column) => (
                    <Table.Td key={column.id}>{column.cell ? column.cell(row) : defaultCell(row, column)}</Table.Td>
                  ))}
                </Table.Row>
              ))}
            </Table.Body>
          </Table>
        </div>
      )}

      {!error && !loading && showPagination ? (
        <Pagination
          page={effectivePage}
          pageCount={pageCount}
          onPageChange={setPage}
          label={pagination?.label ?? `${caption} pagination`}
        />
      ) : null}
    </div>
  )
}