'use client';

import * as React from 'react';

import { ArrowDownIcon, ArrowUpIcon, ChevronsUpDownIcon } from '@/icons';
import { cn } from '@/lib/utils';
import { focusRing } from '@/lib/cva-presets';

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

export interface DataTableSort {
  key: React.Key;
  direction: SortDirection;
}

export interface DataTableColumn<RecordType> {
  /** Required — it identifies the column for sorting and for React keys. */
  key: React.Key;
  title: React.ReactNode;
  /** Path to the value in the record. Omit it and supply `render` instead. */
  dataIndex?: string | number | (string | number)[];
  render?: (value: unknown, record: RecordType, index: number) => React.ReactNode;
  width?: number | string;
  minWidth?: number;
  align?: 'left' | 'center' | 'right';
  /** Pin the column while the rest scrolls horizontally. Needs `scroll.x`. */
  fixed?: 'left' | 'right';
  /** Truncate overflowing text instead of wrapping. */
  ellipsis?: boolean;
  className?: string;
  /** Adds a sort control to this column's header. */
  sortable?: boolean;
  /** Comparator for local sorting. Defaults to a numeric-aware compare. */
  sorter?: (a: RecordType, b: RecordType) => number;
}

export interface DataTableExpandable<RecordType> {
  expandedRowRender?: (record: RecordType, index: number) => React.ReactNode;
  rowExpandable?: (record: RecordType) => boolean;
  defaultExpandAllRows?: boolean;
  expandedRowKeys?: readonly React.Key[];
  onExpand?: (expanded: boolean, record: RecordType) => void;
}

export interface DataTableProps<RecordType> {
  columns: DataTableColumn<RecordType>[];
  data: readonly RecordType[];
  /** Field name or function producing a stable key per row. */
  rowKey?: string | ((record: RecordType, index?: number) => React.Key);
  className?: string;
  rowClassName?: string | ((record: RecordType, index: number) => string);
  /** `{ y }` makes the header sticky; `{ x }` enables horizontal scrolling. */
  scroll?: { x?: number | string | true; y?: number | string };
  /** Keep the header visible while the table's container scrolls. */
  sticky?: boolean;
  expandable?: DataTableExpandable<RecordType>;
  emptyText?: React.ReactNode;
  /**
   * Controlled sort. Pass it together with `onSortChange` when something else
   * sorts the data (a server, a query hook). Leave both off and the table
   * sorts its own rows.
   */
  sort?: DataTableSort | null;
  onSortChange?: (sort: DataTableSort | null) => void;
  onRow?: (record: RecordType, index?: number) => React.HTMLAttributes<HTMLElement>;
  id?: string;
  'aria-label'?: string;
}

/** Reads a column's value out of a record for the default comparator. */
const readCell = <RecordType,>(record: RecordType, column: DataTableColumn<RecordType>) => {
  if (column.dataIndex === undefined) return undefined;
  const segments = Array.isArray(column.dataIndex) ? column.dataIndex : [column.dataIndex];
  return segments.reduce<unknown>(
    (value, segment) =>
      value == null ? value : (value as Record<string | number, unknown>)[segment],
    record
  );
};

const defaultCompare = (a: unknown, b: unknown) => {
  if (a == null && b == null) return 0;
  if (a == null) return -1;
  if (b == null) return 1;
  if (typeof a === 'number' && typeof b === 'number') return a - b;
  return String(a).localeCompare(String(b), undefined, { numeric: true });
};

const alignClass = { left: 'text-left', center: 'text-center', right: 'text-right' } as const;

/** Marker for the extra column that carries the expand toggle. */
const EXPAND_KEY = '__expand__';

/**
 * Table for structured data.
 *
 * A plain `<table>` covers most of it; what this adds is the part that makes a
 * data grid usable at scale — a header that stays put, columns pinned while the
 * rest scrolls sideways, and expandable rows.
 *
 * Sorting works either way round. Omit `sort`/`onSortChange` and the table
 * reorders its own rows; pass both and the header becomes a controlled reporter
 * of intent, which is what you want when a server does the sorting.
 *
 * ```tsx
 * <DataTable
 *   rowKey="id"
 *   data={users}
 *   columns={[
 *     { key: 'name', title: 'Name', dataIndex: 'name', sortable: true },
 *     { key: 'email', title: 'Email', dataIndex: 'email' },
 *   ]}
 * />
 * ```
 */
function DataTable<RecordType extends object>({
  columns,
  data,
  rowKey = 'key',
  className,
  rowClassName,
  scroll,
  sticky,
  expandable,
  emptyText = 'No results.',
  sort: controlledSort,
  onSortChange,
  onRow,
  id,
  ...props
}: DataTableProps<RecordType>) {
  const scrollRef = React.useRef<HTMLDivElement>(null);
  const tableRef = React.useRef<HTMLTableElement>(null);

  const [uncontrolledSort, setUncontrolledSort] = React.useState<DataTableSort | null>(null);
  const isSortControlled = controlledSort !== undefined;
  const sort = isSortControlled ? controlledSort : uncontrolledSort;

  const keyOf = React.useCallback(
    (record: RecordType, index: number): React.Key =>
      typeof rowKey === 'function'
        ? rowKey(record, index)
        : ((record as Record<string, React.Key>)[rowKey] ?? index),
    [rowKey]
  );

  const [uncontrolledExpanded, setUncontrolledExpanded] = React.useState<Set<React.Key>>(() =>
    expandable?.defaultExpandAllRows
      ? new Set(data.map((record, index) => keyOf(record, index)))
      : new Set()
  );

  const expandedKeys = expandable?.expandedRowKeys
    ? new Set(expandable.expandedRowKeys)
    : uncontrolledExpanded;

  const toggleSort = (key: React.Key) => {
    // asc → desc → unsorted, so a mis-click is always one more click from neutral.
    const next: DataTableSort | null =
      sort?.key !== key
        ? { key, direction: 'asc' }
        : sort.direction === 'asc'
          ? { key, direction: 'desc' }
          : null;

    if (!isSortControlled) setUncontrolledSort(next);
    onSortChange?.(next);
  };

  const toggleExpanded = (key: React.Key, record: RecordType) => {
    const willExpand = !expandedKeys.has(key);

    if (!expandable?.expandedRowKeys) {
      setUncontrolledExpanded((current) => {
        const next = new Set(current);
        if (willExpand) next.add(key);
        else next.delete(key);
        return next;
      });
    }

    expandable?.onExpand?.(willExpand, record);
  };

  /* Only sort locally when the caller has not taken control of it. */
  const rows = React.useMemo(() => {
    if (isSortControlled || !sort) return data;

    const column = columns.find((entry) => entry.key === sort.key);
    if (!column) return data;

    const compare =
      column.sorter ??
      ((a: RecordType, b: RecordType) => defaultCompare(readCell(a, column), readCell(b, column)));

    const sorted = [...data].sort(compare);
    return sort.direction === 'desc' ? sorted.reverse() : sorted;
  }, [columns, data, isSortControlled, sort]);

  /* The expand toggle rides in its own leading column, pinned alongside the
     first column when that one is pinned — otherwise it would scroll away from
     the row it belongs to. */
  const leafColumns: DataTableColumn<RecordType>[] = expandable?.expandedRowRender
    ? [
        {
          key: EXPAND_KEY,
          title: '',
          width: 48,
          align: 'center',
          fixed: columns.some((column) => column.fixed === 'left') ? 'left' : undefined,
        },
        ...columns,
      ]
    : columns;

  const lastFixedLeft = leafColumns.reduce(
    (last, column, index) => (column.fixed === 'left' ? index : last),
    -1
  );
  const firstFixedRight = leafColumns.findIndex((column) => column.fixed === 'right');

  /**
   * Pinned columns need a pixel offset that only exists once the browser has
   * laid the table out. Writing it as a custom property straight onto the table
   * keeps it out of React state — there is nothing to re-render, and measuring
   * into state would cost a second pass on every resize and scroll.
   */
  const measure = React.useCallback(() => {
    const table = tableRef.current;
    const viewport = scrollRef.current;
    if (!table || !viewport) return;

    const headers = Array.from(table.querySelectorAll<HTMLElement>('thead th'));

    let left = 0;
    headers.forEach((header, index) => {
      if (header.dataset.fixed !== 'left') return;
      table.style.setProperty(`--ui-table-fixed-${index}`, `${left}px`);
      left += header.offsetWidth;
    });

    let right = 0;
    for (let index = headers.length - 1; index >= 0; index -= 1) {
      if (headers[index].dataset.fixed !== 'right') continue;
      table.style.setProperty(`--ui-table-fixed-${index}`, `${right}px`);
      right += headers[index].offsetWidth;
    }

    /* The shadows say "there is more this way", so they only appear while there
       is somewhere left to scroll. */
    viewport.toggleAttribute('data-ping-left', viewport.scrollLeft > 0);
    viewport.toggleAttribute(
      'data-ping-right',
      Math.ceil(viewport.scrollLeft + viewport.clientWidth) < viewport.scrollWidth
    );
  }, []);

  React.useEffect(() => {
    measure();

    const table = tableRef.current;
    if (!table || typeof ResizeObserver === 'undefined') return;

    const observer = new ResizeObserver(measure);
    observer.observe(table);
    return () => observer.disconnect();
  }, [measure, leafColumns.length, rows]);

  const headerSticky = Boolean(scroll?.y || sticky);

  const pinClasses = (column: DataTableColumn<RecordType>, index: number) =>
    column.fixed
      ? cn(
          'sticky z-20',
          /* The pinned edge fades the content sliding under it — a gradient
             rather than a shadow, so the colour stays a semantic token. */
          'after:pointer-events-none after:absolute after:inset-y-0 after:w-5 after:opacity-0',
          'after:transition-opacity after:duration-(--ui-duration-base)',
          index === lastFixedLeft &&
            'after:left-full after:bg-linear-to-r after:from-overlay/20 after:to-transparent group-data-ping-left/table:after:opacity-100',
          index === firstFixedRight &&
            'after:right-full after:bg-linear-to-l after:from-overlay/20 after:to-transparent group-data-ping-right/table:after:opacity-100'
        )
      : undefined;

  const pinStyle = (column: DataTableColumn<RecordType>, index: number) =>
    column.fixed ? { [column.fixed]: `var(--ui-table-fixed-${index})` } : undefined;

  const cellPadding = 'px-3 py-2.5 align-middle';

  return (
    <div
      ref={scrollRef}
      id={id}
      data-slot="data-table"
      onScroll={measure}
      style={scroll?.y ? { maxHeight: scroll.y } : undefined}
      className={cn(
        'group/table relative rounded-lg border',
        scroll?.x || scroll?.y ? 'overflow-auto' : 'overflow-hidden',
        className
      )}
    >
      <table
        ref={tableRef}
        /* `border-separate` keeps per-cell borders painted under sticky cells,
           which `border-collapse` drops as soon as a cell is lifted out. */
        className="w-full border-separate border-spacing-0 text-sm"
        style={scroll?.x ? { minWidth: scroll.x === true ? '100%' : scroll.x } : undefined}
        {...props}
      >
        <colgroup>
          {leafColumns.map((column) => (
            <col key={column.key} style={column.width ? { width: column.width } : undefined} />
          ))}
        </colgroup>

        <thead>
          <tr>
            {leafColumns.map((column, index) => {
              const active = sort?.key === column.key;
              const SortIcon = !active
                ? ChevronsUpDownIcon
                : sort.direction === 'asc'
                  ? ArrowUpIcon
                  : ArrowDownIcon;

              return (
                <th
                  key={column.key}
                  scope="col"
                  data-fixed={column.fixed}
                  aria-sort={
                    column.sortable
                      ? active
                        ? sort.direction === 'asc'
                          ? 'ascending'
                          : 'descending'
                        : 'none'
                      : undefined
                  }
                  style={{ ...pinStyle(column, index), minWidth: column.minWidth }}
                  className={cn(
                    cellPadding,
                    'border-b bg-muted font-medium whitespace-nowrap text-muted-foreground',
                    headerSticky && 'sticky top-0 z-30',
                    alignClass[column.align ?? 'left'],
                    pinClasses(column, index),
                    column.className
                  )}
                >
                  {column.sortable ? (
                    <button
                      type="button"
                      data-slot="data-table-sort"
                      onClick={() => toggleSort(column.key)}
                      className={cn(
                        '-mx-2 inline-flex items-center gap-1.5 rounded-md px-2 py-1 font-medium',
                        'transition-colors duration-(--ui-duration-fast) hover:text-foreground',
                        active ? 'text-foreground' : 'text-muted-foreground',
                        focusRing
                      )}
                    >
                      {column.title}
                      <SortIcon className={cn('size-3.5', !active && 'opacity-50')} />
                    </button>
                  ) : (
                    column.title
                  )}
                </th>
              );
            })}
          </tr>
        </thead>

        <tbody>
          {rows.length === 0 ? (
            <tr data-slot="data-table-empty">
              <td
                colSpan={leafColumns.length}
                className="px-3 py-8 text-center text-muted-foreground"
              >
                {emptyText}
              </td>
            </tr>
          ) : (
            rows.map((record, index) => {
              const key = keyOf(record, index);
              const renderExpanded = expandable?.expandedRowRender;
              const canExpand = renderExpanded
                ? (expandable.rowExpandable?.(record) ?? true)
                : false;
              const expanded = expandedKeys.has(key);
              const rowProps = onRow?.(record, index);

              return (
                <React.Fragment key={key}>
                  <tr
                    data-slot="data-table-row"
                    {...rowProps}
                    className={cn(
                      'group/row',
                      typeof rowClassName === 'function'
                        ? rowClassName(record, index)
                        : rowClassName,
                      rowProps?.className
                    )}
                  >
                    {leafColumns.map((column, columnIndex) => {
                      const isExpandCell = column.key === EXPAND_KEY;
                      const value = readCell(record, column);

                      return (
                        <td
                          key={column.key}
                          style={pinStyle(column, columnIndex)}
                          className={cn(
                            cellPadding,
                            /* Pinned cells carry their own background so they
                               stay opaque over the rows sliding underneath. */
                            'border-b bg-background group-hover/row:bg-muted/60',
                            alignClass[column.align ?? 'left'],
                            column.ellipsis && 'truncate',
                            pinClasses(column, columnIndex),
                            column.className
                          )}
                        >
                          {isExpandCell ? (
                            canExpand ? (
                              <button
                                type="button"
                                aria-expanded={expanded}
                                aria-label={expanded ? 'Collapse row' : 'Expand row'}
                                data-slot="data-table-expand"
                                onClick={() => toggleExpanded(key, record)}
                                className={cn(
                                  'inline-flex size-[1.125rem] items-center justify-center rounded-sm border text-muted-foreground',
                                  'transition-colors duration-(--ui-duration-fast) hover:bg-accent hover:text-accent-foreground',
                                  focusRing
                                )}
                              >
                                {expanded ? '−' : '+'}
                              </button>
                            ) : null
                          ) : column.render ? (
                            column.render(value, record, index)
                          ) : (
                            (value as React.ReactNode)
                          )}
                        </td>
                      );
                    })}
                  </tr>

                  {renderExpanded && expanded && canExpand ? (
                    <tr data-slot="data-table-expanded-row">
                      <td
                        colSpan={leafColumns.length}
                        className={cn(cellPadding, 'border-b bg-muted/45')}
                      >
                        {renderExpanded(record, index)}
                      </td>
                    </tr>
                  ) : null}
                </React.Fragment>
              );
            })
          )}
        </tbody>
      </table>
    </div>
  );
}

export { DataTable };