import { useMemo, useRef, useState, useCallback } from "react";
import {
  useReactTable,
  getCoreRowModel,
  getSortedRowModel,
  flexRender,
  type ColumnDef,
  type SortingState,
  type ColumnResizeMode,
} from "@tanstack/react-table";
import { useVirtualizer } from "@tanstack/react-virtual";
import { ArrowUp, ArrowDown } from "lucide-react";

interface Props {
  columns: string[];
  types: string[];
  rows: unknown[][];
  onCellDoubleClick?: (
    rowIndex: number,
    colIndex: number,
    value: unknown
  ) => void;
  highlightedCells?: Set<string>;
}

export function ResultsTable({
  columns: colNames,
  rows,
  onCellDoubleClick,
  highlightedCells,
}: Props) {
  const [sorting, setSorting] = useState<SortingState>([]);
  const [columnResizeMode] = useState<ColumnResizeMode>("onChange");
  const parentRef = useRef<HTMLDivElement>(null);

  const columns = useMemo<ColumnDef<unknown[]>[]>(
    () =>
      colNames.map((name, i) => ({
        id: name,
        accessorFn: (row: unknown[]) => row[i],
        header: name,
        cell: ({ getValue, row: tableRow }) => {
          const value = getValue();
          const rowIndex = tableRow.index;
          const key = `${rowIndex}:${i}`;
          const isHighlighted = highlightedCells?.has(key);

          return (
            <div
              className={`truncate px-2 py-1 ${
                value === null
                  ? "italic text-muted-foreground"
                  : isHighlighted
                    ? "bg-yellow-900/30"
                    : ""
              }`}
              onDoubleClick={() =>
                onCellDoubleClick?.(rowIndex, i, value)
              }
            >
              {value === null ? "NULL" : String(value)}
            </div>
          );
        },
        size: 150,
        minSize: 50,
        maxSize: 800,
      })),
    [colNames, onCellDoubleClick, highlightedCells]
  );

  const table = useReactTable({
    data: rows,
    columns,
    state: { sorting },
    onSortingChange: setSorting,
    getCoreRowModel: getCoreRowModel(),
    getSortedRowModel: getSortedRowModel(),
    columnResizeMode,
    enableColumnResizing: true,
  });

  const { rows: tableRows } = table.getRowModel();

  const virtualizer = useVirtualizer({
    count: tableRows.length,
    getScrollElement: () => parentRef.current,
    estimateSize: () => 28,
    overscan: 20,
  });

  const columnSizing = table.getState().columnSizing;

  const getColWidth = useCallback(
    (colId: string, fallback: number) => columnSizing[colId] ?? fallback,
    [columnSizing]
  );

  if (colNames.length === 0) return null;

  return (
    <div ref={parentRef} className="h-full overflow-auto">
      {/* Header */}
      <div className="sticky top-0 z-10 flex bg-card border-b">
        {table.getHeaderGroups().map((headerGroup) =>
          headerGroup.headers.map((header) => (
            <div
              key={header.id}
              className="relative shrink-0 select-none border-r px-2 py-1.5 text-left text-xs font-medium text-muted-foreground"
              style={{ width: header.getSize(), minWidth: header.getSize() }}
            >
              <div
                className="flex cursor-pointer items-center gap-1 hover:text-foreground"
                onClick={header.column.getToggleSortingHandler()}
              >
                {flexRender(
                  header.column.columnDef.header,
                  header.getContext()
                )}
                {header.column.getIsSorted() === "asc" && (
                  <ArrowUp className="h-3 w-3" />
                )}
                {header.column.getIsSorted() === "desc" && (
                  <ArrowDown className="h-3 w-3" />
                )}
              </div>
              {/* Resize handle */}
              <div
                onMouseDown={header.getResizeHandler()}
                onTouchStart={header.getResizeHandler()}
                onDoubleClick={() => header.column.resetSize()}
                className={`absolute right-0 top-0 h-full w-1 cursor-col-resize select-none touch-none hover:bg-primary ${
                  header.column.getIsResizing() ? "bg-primary" : ""
                }`}
              />
            </div>
          ))
        )}
      </div>

      {/* Virtual rows */}
      <div
        style={{
          height: `${virtualizer.getTotalSize()}px`,
          position: "relative",
        }}
      >
        {virtualizer.getVirtualItems().map((virtualRow) => {
          const row = tableRows[virtualRow.index];
          return (
            <div
              key={row.id}
              className="absolute left-0 flex hover:bg-accent/50"
              style={{
                top: `${virtualRow.start}px`,
                height: `${virtualRow.size}px`,
              }}
            >
              {row.getVisibleCells().map((cell) => {
                const w = getColWidth(cell.column.id, cell.column.getSize());
                return (
                  <div
                    key={cell.id}
                    className="shrink-0 border-b border-r text-xs"
                    style={{ width: w, minWidth: w }}
                  >
                    {flexRender(
                      cell.column.columnDef.cell,
                      cell.getContext()
                    )}
                  </div>
                );
              })}
            </div>
          );
        })}
      </div>
    </div>
  );
}