import { cn } from '@learnitive/ui-components';
import React from 'react';

export type LatexTableGridSize = {
  rows: number;
  columns: number;
};

export function LatexTableGridPicker({
  rows = 8,
  columns = 8,
  onSelect,
  onCancel,
  className,
}: {
  rows?: number;
  columns?: number;
  onSelect: (size: LatexTableGridSize) => void;
  onCancel?: () => void;
  className?: string;
}) {
  const rowCount = Math.max(1, Math.min(16, Math.trunc(rows)));
  const columnCount = Math.max(1, Math.min(16, Math.trunc(columns)));
  const [active, setActive] = React.useState<LatexTableGridSize>();
  const buttonsRef = React.useRef<Array<HTMLButtonElement | null>>([]);

  const focusCell = React.useCallback((nextRows: number, nextColumns: number) => {
    const next = {
      rows: Math.max(1, Math.min(rowCount, nextRows)),
      columns: Math.max(1, Math.min(columnCount, nextColumns)),
    };
    setActive(next);
    const index = (next.rows - 1) * columnCount + (next.columns - 1);
    buttonsRef.current[index]?.focus();
  }, [columnCount, rowCount]);

  const handleKeyDown = (event: React.KeyboardEvent<HTMLButtonElement>, size: LatexTableGridSize) => {
    if (event.key === 'Escape') {
      event.preventDefault();
      onCancel?.();
      return;
    }
    if (event.key === 'Enter' || event.key === ' ') {
      event.preventDefault();
      onSelect(size);
      return;
    }
    if (event.key === 'ArrowLeft') {
      event.preventDefault();
      focusCell(size.rows, size.columns - 1);
    } else if (event.key === 'ArrowRight') {
      event.preventDefault();
      focusCell(size.rows, size.columns + 1);
    } else if (event.key === 'ArrowUp') {
      event.preventDefault();
      focusCell(size.rows - 1, size.columns);
    } else if (event.key === 'ArrowDown') {
      event.preventDefault();
      focusCell(size.rows + 1, size.columns);
    }
  };

  return (
    <div
      className={cn('grid gap-2 p-1', className)}
      data-testid="latex-table-grid-picker"
    >
      <div className="flex min-h-5 items-center justify-between gap-4 text-xs">
        <span className="font-medium">Select size</span>
        <span className="min-w-12 text-right tabular-nums text-muted-foreground" aria-live="polite">
          {active ? `${active.rows} × ${active.columns}` : ''}
        </span>
      </div>
      <div
        className="grid gap-1"
        style={{ gridTemplateColumns: `repeat(${columnCount}, minmax(0, 1fr))` }}
        role="grid"
        onMouseLeave={() => setActive(undefined)}
        aria-label={`Select table size, up to ${rowCount} by ${columnCount}`}
      >
        {Array.from({ length: rowCount }, (_, rowIndex) =>
          Array.from({ length: columnCount }, (_, columnIndex) => {
            const size = { rows: rowIndex + 1, columns: columnIndex + 1 };
            const selected = Boolean(
              active && size.rows <= active.rows && size.columns <= active.columns,
            );
            const index = rowIndex * columnCount + columnIndex;
            return (
              <button
                key={`${size.rows}:${size.columns}`}
                ref={(node) => {
                  buttonsRef.current[index] = node;
                }}
                type="button"
                role="gridcell"
                aria-label={`${size.rows} by ${size.columns} table`}
                aria-selected={selected}
                data-table-grid-row={size.rows}
                data-table-grid-column={size.columns}
                className={cn(
                  'size-6 rounded-sm border transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring',
                  selected ? 'border-primary bg-primary/20' : 'border-border bg-background hover:bg-accent',
                )}
                onMouseEnter={() => setActive(size)}
                onFocus={() => setActive(size)}
                onKeyDown={(event) => handleKeyDown(event, size)}
                onClick={() => onSelect(size)}
              />
            );
          }),
        )}
      </div>
    </div>
  );
}