oranguru/src/Griddy/core/Griddy.tsx

import {
  type ColumnDef,
  type ColumnFiltersState,
  type ColumnOrderState,
  type ColumnPinningState,
  getCoreRowModel,
  getExpandedRowModel,
  getFilteredRowModel,
  getGroupedRowModel,
  getPaginationRowModel,
  getSortedRowModel,
  type GroupingState,
  type PaginationState,
  type RowSelectionState,
  type SortingState,
  useReactTable,
  type VisibilityState,
} from '@tanstack/react-table';
import React, {
  forwardRef,
  type Ref,
  useEffect,
  useImperativeHandle,
  useMemo,
  useRef,
  useState,
} from 'react';

import type { GriddyProps, GriddyRef } from './types';

import { advancedSearchGlobalFilterFn, AdvancedSearchPanel } from '../features/advancedSearch';
import { GriddyErrorBoundary } from '../features/errorBoundary';
import { useKeyboardNavigation } from '../features/keyboard/useKeyboardNavigation';
import { GriddyLoadingOverlay, GriddyLoadingSkeleton } from '../features/loading';
import { PaginationControl } from '../features/pagination';
import { SearchOverlay } from '../features/search/SearchOverlay';
import { GridToolbar } from '../features/toolbar';
import { useGridVirtualizer } from '../rendering/hooks/useGridVirtualizer';
import { TableHeader } from '../rendering/TableHeader';
import { VirtualBody } from '../rendering/VirtualBody';
import styles from '../styles/griddy.module.css';
import { mapColumns } from './columnMapper';
import { CSS, DEFAULTS } from './constants';
import { GriddyProvider, useGriddyStore } from './GriddyStore';

// ─── Inner Component (lives inside Provider, has store access) ───────────────

function _Griddy<T>(props: GriddyProps<T>, ref: Ref<GriddyRef<T>>) {
  return (
    <GriddyProvider {...props}>
      <GriddyErrorBoundary onError={props.onError} onRetry={props.onRetry}>
        <GriddyInner tableRef={ref} />
      </GriddyErrorBoundary>
      {props.children}
    </GriddyProvider>
  );
}

// ─── Main Component with forwardRef ──────────────────────────────────────────

function GriddyInner<T>({ tableRef }: { tableRef: Ref<GriddyRef<T>> }) {
  // Read props from synced store
  const data = useGriddyStore((s) => s.data);
  const userColumns = useGriddyStore((s) => s.columns);
  const getRowId = useGriddyStore((s) => s.getRowId);
  const selection = useGriddyStore((s) => s.selection);
  const search = useGriddyStore((s) => s.search);
  const groupingConfig = useGriddyStore((s) => s.grouping);
  const paginationConfig = useGriddyStore((s) => s.pagination);
  const controlledSorting = useGriddyStore((s) => s.sorting);
  const onSortingChange = useGriddyStore((s) => s.onSortingChange);
  const controlledFilters = useGriddyStore((s) => s.columnFilters);
  const onColumnFiltersChange = useGriddyStore((s) => s.onColumnFiltersChange);
  const controlledPinning = useGriddyStore((s) => s.columnPinning);
  const onColumnPinningChange = useGriddyStore((s) => s.onColumnPinningChange);
  const controlledRowSelection = useGriddyStore((s) => s.rowSelection);
  const onRowSelectionChange = useGriddyStore((s) => s.onRowSelectionChange);
  const onEditCommit = useGriddyStore((s) => s.onEditCommit);
  const rowHeight = useGriddyStore((s) => s.rowHeight);
  const overscanProp = useGriddyStore((s) => s.overscan);
  const height = useGriddyStore((s) => s.height);
  const keyboardNavigation = useGriddyStore((s) => s.keyboardNavigation);
  const className = useGriddyStore((s) => s.className);
  const showToolbar = useGriddyStore((s) => s.showToolbar);
  const exportFilename = useGriddyStore((s) => s.exportFilename);
  const isLoading = useGriddyStore((s) => s.isLoading);
  const filterPresets = useGriddyStore((s) => s.filterPresets);
  const advancedSearch = useGriddyStore((s) => s.advancedSearch);
  const persistenceKey = useGriddyStore((s) => s.persistenceKey);
  const manualSorting = useGriddyStore((s) => s.manualSorting);
  const manualFiltering = useGriddyStore((s) => s.manualFiltering);
  const dataCount = useGriddyStore((s) => s.dataCount);
  const setTable = useGriddyStore((s) => s.setTable);
  const setVirtualizer = useGriddyStore((s) => s.setVirtualizer);
  const setScrollRef = useGriddyStore((s) => s.setScrollRef);
  const setFocusedRow = useGriddyStore((s) => s.setFocusedRow);
  const setFocusedColumn = useGriddyStore((s) => s.setFocusedColumn);
  const setEditing = useGriddyStore((s) => s.setEditing);
  const setTotalRows = useGriddyStore((s) => s.setTotalRows);
  const focusedRowIndex = useGriddyStore((s) => s.focusedRowIndex);

  const effectiveRowHeight = rowHeight ?? DEFAULTS.rowHeight;
  const effectiveOverscan = overscanProp ?? DEFAULTS.overscan;
  const enableKeyboard = keyboardNavigation !== false;

  // ─── Column Mapping ───
  const columns = useMemo(
    () => mapColumns(userColumns ?? [], selection) as ColumnDef<T, any>[],
    [userColumns, selection]
  );

  // ─── Table State (internal/uncontrolled) ───
  const [internalSorting, setInternalSorting] = useState<SortingState>([]);
  const [internalFilters, setInternalFilters] = useState<ColumnFiltersState>([]);
  const [internalRowSelection, setInternalRowSelection] = useState<RowSelectionState>({});
  const [globalFilter, setGlobalFilter] = useState<string | undefined>(undefined);
  const [columnVisibility, setColumnVisibility] = useState<VisibilityState>({});
  const [columnOrder, setColumnOrder] = useState<ColumnOrderState>([]);

  // Build initial column pinning from column definitions
  const initialPinning = useMemo(() => {
    const left: string[] = [];
    const right: string[] = [];
    userColumns?.forEach((col) => {
      if (col.pinned === 'left') left.push(col.id);
      else if (col.pinned === 'right') right.push(col.id);
    });
    return { left, right };
  }, [userColumns]);

  const [internalPinning, setInternalPinning] = useState<ColumnPinningState>(initialPinning);
  const [grouping, setGrouping] = useState<GroupingState>(groupingConfig?.columns ?? []);
  const [expanded, setExpanded] = useState({});
  const [internalPagination, setInternalPagination] = useState<PaginationState>({
    pageIndex: 0,
    pageSize: paginationConfig?.pageSize ?? DEFAULTS.pageSize,
  });

  // Wrap pagination setters to call callbacks
  const handlePaginationChange = (updater: any) => {
    setInternalPagination((prev) => {
      const next = typeof updater === 'function' ? updater(prev) : updater;
      // Call callbacks if pagination config exists
      if (paginationConfig) {
        if (next.pageIndex !== prev.pageIndex && paginationConfig.onPageChange) {
          paginationConfig.onPageChange(next.pageIndex);
        }
        if (next.pageSize !== prev.pageSize && paginationConfig.onPageSizeChange) {
          paginationConfig.onPageSizeChange(next.pageSize);
        }
      }
      return next;
    });
  };

  // Resolve controlled vs uncontrolled
  const sorting = controlledSorting ?? internalSorting;
  const setSorting = onSortingChange ?? setInternalSorting;
  const columnFilters = controlledFilters ?? internalFilters;
  const setColumnFilters = onColumnFiltersChange ?? setInternalFilters;
  const columnPinning = controlledPinning ?? internalPinning;
  const setColumnPinning = onColumnPinningChange ?? setInternalPinning;
  const rowSelectionState = controlledRowSelection ?? internalRowSelection;
  const setRowSelection = onRowSelectionChange ?? setInternalRowSelection;

  // ─── Selection config ───
  const enableRowSelection = selection ? selection.mode !== 'none' : false;
  const enableMultiRowSelection = selection?.mode === 'multi';

  // ─── TanStack Table Instance ───
  const table = useReactTable<T>({
    columns,
    data: (data ?? []) as T[],
    enableColumnResizing: true,
    enableExpanding: true,
    enableFilters: true,
    enableGrouping: groupingConfig?.enabled ?? false,
    enableMultiRowSelection,
    enableMultiSort: true,
    enablePinning: true,
    enableRowSelection,
    enableSorting: true,
    getCoreRowModel: getCoreRowModel(),
    ...(advancedSearch?.enabled ? { globalFilterFn: advancedSearchGlobalFilterFn as any } : {}),
    getExpandedRowModel: getExpandedRowModel(),
    getFilteredRowModel: manualFiltering ? undefined : getFilteredRowModel(),
    getGroupedRowModel: groupingConfig?.enabled ? getGroupedRowModel() : undefined,
    getRowId: (getRowId as any) ?? ((_, index) => String(index)),
    getSortedRowModel: manualSorting ? undefined : getSortedRowModel(),
    manualFiltering: manualFiltering ?? false,
    manualSorting: manualSorting ?? false,
    onColumnFiltersChange: setColumnFilters as any,
    onColumnOrderChange: setColumnOrder,
    onColumnPinningChange: setColumnPinning as any,
    onColumnVisibilityChange: setColumnVisibility,
    onExpandedChange: setExpanded,
    onGlobalFilterChange: setGlobalFilter,
    onGroupingChange: setGrouping,
    onPaginationChange: paginationConfig?.enabled ? handlePaginationChange : undefined,
    onRowSelectionChange: setRowSelection as any,
    onSortingChange: setSorting as any,
    rowCount: dataCount,
    state: {
      columnFilters,
      columnOrder,
      columnPinning,
      columnVisibility,
      expanded,
      globalFilter,
      grouping,
      rowSelection: rowSelectionState,
      sorting,
      ...(paginationConfig?.enabled ? { pagination: internalPagination } : {}),
    },
    ...(paginationConfig?.enabled ? { getPaginationRowModel: getPaginationRowModel() } : {}),
    columnResizeMode: 'onChange',
  });

  // ─── Scroll Container Ref ───
  const scrollRef = useRef<HTMLDivElement>(null);

  // ─── TanStack Virtual ───
  const virtualizer = useGridVirtualizer({
    overscan: effectiveOverscan,
    rowHeight: effectiveRowHeight,
    scrollRef,
    table,
  });

  // ─── Sync table + virtualizer + scrollRef into store ───
  useEffect(() => {
    setTable(table);
  }, [table, setTable]);
  useEffect(() => {
    setVirtualizer(virtualizer);
  }, [virtualizer, setVirtualizer]);
  useEffect(() => {
    setScrollRef(scrollRef.current);
  }, [setScrollRef]);

  // ─── Keyboard Navigation ───
  // Get the full store state for imperative access in keyboard handler
  const storeState = useGriddyStore();

  useKeyboardNavigation({
    editingEnabled: !!onEditCommit,
    scrollRef,
    search,
    selection,
    storeState,
    table,
    virtualizer,
  });

  // ─── Set initial focus when data loads ───
  const rowCount = table.getRowModel().rows.length;

  useEffect(() => {
    setTotalRows(rowCount);
    if (rowCount > 0 && focusedRowIndex === null) {
      setFocusedRow(0);
    }
  }, [rowCount, focusedRowIndex, setTotalRows, setFocusedRow]);

  // ─── Imperative Ref ───
  useImperativeHandle(
    tableRef,
    () => ({
      deselectAll: () => table.resetRowSelection(),
      focusRow: (index: number) => {
        setFocusedRow(index);
        virtualizer.scrollToIndex(index, { align: 'auto' });
      },
      getTable: () => table,
      getUIState: () =>
        ({
          focusedColumnId: null,
          focusedRowIndex,
          isEditing: false,
          isSearchOpen: false,
          isSelecting: false,
          totalRows: rowCount,
        }) as any,
      getVirtualizer: () => virtualizer,
      scrollToRow: (index: number) => virtualizer.scrollToIndex(index, { align: 'auto' }),
      selectRow: (id: string) => {
        const row = table.getRowModel().rows.find((r) => r.id === id);
        row?.toggleSelected(true);
      },
      startEditing: (rowId: string, columnId?: string) => {
        const rowIndex = table.getRowModel().rows.findIndex((r) => r.id === rowId);
        if (rowIndex >= 0) {
          setFocusedRow(rowIndex);
          if (columnId) setFocusedColumn(columnId);
          setEditing(true);
        }
      },
    }),
    [table, virtualizer, setFocusedRow, setFocusedColumn, setEditing, focusedRowIndex, rowCount]
  );

  // ─── Render ───
  const containerStyle: React.CSSProperties = {
    height: height ?? '100%',
    overflow: 'auto',
    position: 'relative',
  };

  const focusedRow = focusedRowIndex !== null ? table.getRowModel().rows[focusedRowIndex] : null;
  const focusedRowId = focusedRow ? `griddy-row-${focusedRow.id}` : undefined;

  return (
    <div
      aria-activedescendant={focusedRowId}
      aria-label="Data grid"
      aria-rowcount={(data ?? []).length}
      className={[styles[CSS.root], className].filter(Boolean).join(' ')}
      role="grid"
    >
      {search?.enabled && <SearchOverlay />}
      {advancedSearch?.enabled && <AdvancedSearchPanel table={table} />}
      {showToolbar && (
        <GridToolbar
          exportFilename={exportFilename}
          filterPresets={filterPresets}
          persistenceKey={persistenceKey}
          table={table}
        />
      )}
      <div
        className={styles[CSS.container]}
        ref={scrollRef}
        style={containerStyle}
        tabIndex={enableKeyboard ? 0 : undefined}
      >
        <TableHeader />
        {isLoading && (!data || data.length === 0) ? (
          <GriddyLoadingSkeleton />
        ) : (
          <>
            <VirtualBody />
            {isLoading && <GriddyLoadingOverlay />}
          </>
        )}
      </div>
      {paginationConfig?.enabled && (
        <PaginationControl pageSizeOptions={paginationConfig.pageSizeOptions} table={table} />
      )}
    </div>
  );
}

export const Griddy = forwardRef(_Griddy) as <T>(
  props: GriddyProps<T> & React.RefAttributes<GriddyRef<T>>
) => React.ReactElement;