useTable.tsx

'use client';

import React, { CSSProperties, useState } from 'react'

import {
  Cell,
  ColumnDef,
  ColumnFiltersState,
  Header,
  SortingState,
  VisibilityState,
  ExpandedState,
  flexRender,
  getCoreRowModel,
  getFilteredRowModel,
  getPaginationRowModel,
  getSortedRowModel,
  useReactTable,
  OnChangeFn,
  Row,
  getExpandedRowModel,
} from '@tanstack/react-table'

// needed for table body level scope DnD setup
import {
  DndContext,
  KeyboardSensor,
  MouseSensor,
  TouchSensor,
  closestCenter,
  type DragEndEvent,
  useSensor,
  useSensors,
} from '@dnd-kit/core'
import { restrictToHorizontalAxis } from '@dnd-kit/modifiers'
import {
  arrayMove,
  SortableContext,
  horizontalListSortingStrategy,
} from '@dnd-kit/sortable'

// needed for row & cell level scope DnD setup
import { useSortable } from '@dnd-kit/sortable'
import { CSS } from '@dnd-kit/utilities'



export interface TableConfig<TData, TLevelData> {
  columns: ColumnDef<TData>[];
  data: TData[];
  enableSorting?: boolean;
  enableFiltering?: boolean;
  enableColumnVisibility?: boolean;
  enableDragAndDrop?: boolean;
  enablePagination?: boolean;
  columnsByLevel?: Record<number, ColumnDef<TData>[]>;
  maxDepth?: number;
  subRowsKey?: string | Record<number, string>;
}


const DraggableTableHeader = ({
  header,
  enableSorting
}: {
  header: Header<any, any>,
  enableSorting: boolean
}) => {
  const { attributes, isDragging, listeners, setNodeRef, transform } =
    useSortable({
      id: header.column.id,
    })

  const style: CSSProperties = {
    opacity: isDragging ? 0.8 : 1,
    transform: CSS.Translate.toString(transform),
    transition: 'width transform 0.2s ease-in-out',
    whiteSpace: 'nowrap',
    width: header.column.getSize(),
    cursor: 'pointer',
    zIndex: isDragging ? 1 : 0,
  }

  return (
    <th
      colSpan={header.colSpan}
      ref={setNodeRef}
      className='flex items-center justify-between px-4 py-2'
      style={style}
      onClick={header.column.getToggleSortingHandler()}
    >
      <div {...attributes} {...listeners} className="flex items-center gap-2">
        {header.isPlaceholder
          ? null
          : flexRender(header.column.columnDef.header, header.getContext())}
      </div>
      {enableSorting && (
        <span className="text-gray-400">
          {{ asc: '↑', desc: '↓' }[header.column.getIsSorted() as string] ?? '↕'}
        </span>
      )}
    </th>
  )
}

const DragAlongCell = ({ cell }: { cell: Cell<any, any> }) => {
  const { isDragging, setNodeRef, transform } = useSortable({
    id: cell.column.id,
  })

  const style: CSSProperties = {
    opacity: isDragging ? 0.8 : 1,
    transform: CSS.Translate.toString(transform), // translate instead of transform to avoid squishing
    transition: 'width transform 0.2s ease-in-out',
    width: cell.column.getSize(),
    zIndex: isDragging ? 1 : 0,
  }

  return (
    <td style={style} ref={setNodeRef}>
      {flexRender(cell.column.columnDef.cell, cell.getContext())}
    </td>
  )
}

// Separate component for nested tables
const NestedTable = ({
  data,
  columns,
  columnsByLevel,
  depth,
  columnFilters,
  columnVisibility,
  enableSorting,
  enableFiltering,
  enablePagination,
  onColumnFiltersChange,
  onColumnVisibilityChange,
  subRowsKey = 'subRows',
}: {
  data: any[],
  columns: ColumnDef<any, any>[],
  columnsByLevel?: Record<number, ColumnDef<any, any>[]>,
  depth: number,
  columnFilters: ColumnFiltersState,
  columnVisibility: VisibilityState,
  enableSorting: boolean,
  enableFiltering: boolean,
  enablePagination: boolean,
  onColumnFiltersChange: OnChangeFn<ColumnFiltersState>,
  onColumnVisibilityChange: OnChangeFn<VisibilityState>,
  subRowsKey?: string | Record<number, string>,
}) => {
  const [sorting, setSorting] = useState<SortingState>([]);
  const [expanded, setExpanded] = useState<ExpandedState>({});

  // Use level-specific columns if available, otherwise use default columns
  const currentLevelColumns = columnsByLevel?.[depth] || columns;

  // Get the current level's subRowsKey
  const getCurrentSubRowsKey = (row: any) => {
    if (typeof subRowsKey === 'string') {
      return row[subRowsKey];
    }
    // Try the current depth's key first
    const currentKey = subRowsKey?.[depth];
    if (currentKey && row[currentKey]) {
      return row[currentKey];
    }
    // If not found, try the next level's key
    const nextKey = subRowsKey?.[depth + 1];
    if (nextKey && row[nextKey]) {
      return row[nextKey];
    }
    // Fallback to default
    return row.subRows;
  };

  const [columnOrder, setColumnOrder] = React.useState<string[]>(() =>
    currentLevelColumns.map(c => c.id!)
  );

  const tableInstance = useReactTable<any>({
    data,
    columns: currentLevelColumns,
    state: {
      sorting,
      columnOrder,
      columnFilters,
      columnVisibility,
      expanded,
    },
    onExpandedChange: setExpanded,
    getSubRows: getCurrentSubRowsKey,
    onColumnOrderChange: setColumnOrder,
    onSortingChange: setSorting,
    onColumnFiltersChange,
    onColumnVisibilityChange,
    getCoreRowModel: getCoreRowModel(),
    getSortedRowModel: enableSorting ? getSortedRowModel() : undefined,
    getFilteredRowModel: enableFiltering ? getFilteredRowModel() : undefined,
    getPaginationRowModel: enablePagination ? getPaginationRowModel() : undefined,
  });

  const sensors = useSensors(
    useSensor(MouseSensor, {}),
    useSensor(TouchSensor, {}),
    useSensor(KeyboardSensor, {})
  );

  function handleDragEnd(event: DragEndEvent) {
    const { active, over } = event;
    if (active && over && active.id !== over.id) {
      const oldIndex = columnOrder.indexOf(active.id as string);
      const newIndex = columnOrder.indexOf(over.id as string);
      setColumnOrder(arrayMove(columnOrder, oldIndex, newIndex));
    }
  }

  return (
    <DndContext
      collisionDetection={closestCenter}
      modifiers={[restrictToHorizontalAxis]}
      onDragEnd={handleDragEnd}
      sensors={sensors}
    >
      <table className="w-full border-collapse">
        <thead>
          {tableInstance.getHeaderGroups().map(headerGroup => (
            <tr key={headerGroup.id} className='flex items-center gap-8'>
              <th className="w-10"></th>
              <SortableContext
                items={columnOrder}
                strategy={horizontalListSortingStrategy}
              >
                {headerGroup.headers.map(header => (
                  <DraggableTableHeader key={header.id} header={header as Header<any, any>} enableSorting={enableSorting} />
                ))}
              </SortableContext>
            </tr>
          ))}
        </thead>
        <tbody>
          {tableInstance.getRowModel().rows.map(row => (
            <React.Fragment key={row.id}>
              <tr className='flex items-center gap-8 hover:bg-gray-50'>
                <td className="w-10">
                  {row.getCanExpand() ? (
                    <button
                      onClick={row.getToggleExpandedHandler()}
                      className="w-6 h-6 flex items-center justify-center"
                    >
                      {row.getIsExpanded() ? '▼' : '▶'}
                    </button>
                  ) : null}
                </td>
                {row.getVisibleCells().map(cell => (
                  <SortableContext
                    key={cell.id}
                    items={columnOrder}
                    strategy={horizontalListSortingStrategy}
                  >
                    <DragAlongCell key={cell.id} cell={cell as Cell<any, any>} />
                  </SortableContext>
                ))}
              </tr>
              {row.getIsExpanded() && getCurrentSubRowsKey(row.original)?.length > 0 && (
                <tr>
                  <td colSpan={row.getVisibleCells().length + 1}>
                    <div className={`pl-8 border-l-2 border-gray-200 my-2`}>
                      <NestedTable
                        data={getCurrentSubRowsKey(row.original)}
                        columns={columns}
                        columnsByLevel={columnsByLevel}
                        depth={depth + 1}
                        columnFilters={columnFilters}
                        columnVisibility={columnVisibility}
                        enableSorting={enableSorting}
                        enableFiltering={enableFiltering}
                        enablePagination={enablePagination}
                        onColumnFiltersChange={onColumnFiltersChange}
                        onColumnVisibilityChange={onColumnVisibilityChange}
                        subRowsKey={subRowsKey}
                      />
                    </div>
                  </td>
                </tr>
              )}
            </React.Fragment>
          ))}
        </tbody>
      </table>
    </DndContext>
  );
};

export function useTable<TData, TLevelData>({
  columns,
  data,
  columnsByLevel,
  maxDepth = 2,
  subRowsKey = 'subRows'
}: TableConfig<TData, TLevelData>) {
  const [columnFilters, setColumnFilters] = useState<ColumnFiltersState>([]);
  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({});
  const [rows, setRows] = useState<TData[]>(data);

  const getSubRows = (row: TData, level: number): TData[] | undefined => {
    const key = typeof subRowsKey === 'string'
      ? subRowsKey
      : (subRowsKey?.[level] || 'subRows');
    return (row as any)[key];
  };

  const getRowCanExpand = (row: Row<TData>): boolean => {
    const subRows = getSubRows(row.original, row.depth);
    return Boolean(subRows && subRows.length);
  };

  const render = () => (
    <div className="p-2">
      <div className="h-4 m-20" />
      <div className="h-4" />
      <NestedTable
        data={data}
        columns={columns}
        columnsByLevel={columnsByLevel}
        depth={0}
        columnFilters={columnFilters}
        columnVisibility={columnVisibility}
        enableSorting={true}
        enableFiltering={false}
        enablePagination={false}
        onColumnFiltersChange={setColumnFilters}
        onColumnVisibilityChange={setColumnVisibility}
        subRowsKey={subRowsKey}
      />
    </div>
  );

  return { render };
}