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 ExternalSort {
  column: string | undefined;
  direction: string | undefined;
  onSort: (column: string | undefined, direction: string | undefined) => void;
}

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

export function ResultsTable({
  columns: colNames,
  rows,
  onCellDoubleClick,
  highlightedCells,
  externalSort,
}: 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 font-mono text-[12px] ${
                value === null
                  ? "tusk-grid-cell-null"
                  : isHighlighted
                    ? "tusk-grid-cell-highlight"
                    : ""
              }`}
              onDoubleClick={() =>
                onCellDoubleClick?.(rowIndex, i, value)
              }
            >
              {value === null ? "NULL" : String(value)}
            </div>
          );
        },
        size: 150,
        minSize: 50,
        maxSize: 800,
      })),
    [colNames, onCellDoubleClick, highlightedCells]
  );

  // eslint-disable-next-line react-hooks/incompatible-library
  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]
  );

  const handleHeaderClick = useCallback(
    (colName: string, defaultHandler: ((e: unknown) => void) | undefined) =>
      (e: unknown) => {
        if (externalSort) {
          if (externalSort.column !== colName) {
            externalSort.onSort(colName, "ASC");
          } else if (externalSort.direction === "ASC") {
            externalSort.onSort(colName, "DESC");
          } else {
            externalSort.onSort(undefined, undefined);
          }
        } else {
          defaultHandler?.(e);
        }
      },
    [externalSort]
  );

  const getIsSorted = useCallback(
    (colName: string, localSorted: false | "asc" | "desc") => {
      if (externalSort) {
        if (externalSort.column === colName) {
          return externalSort.direction === "ASC" ? "asc" : externalSort.direction === "DESC" ? "desc" : false;
        }
        return false;
      }
      return localSorted;
    },
    [externalSort]
  );

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

  return (
    <div ref={parentRef} className="h-full select-text overflow-auto">
      {/* Header */}
      <div className="tusk-grid-header sticky top-0 z-10 flex">
        {table.getHeaderGroups().map((headerGroup) =>
          headerGroup.headers.map((header) => (
            <div
              key={header.id}
              className="relative shrink-0 select-none border-r border-border/30 px-2 py-1.5 text-left text-[11px] font-semibold tracking-wide text-muted-foreground uppercase"
              style={{ width: header.getSize(), minWidth: header.getSize() }}
            >
              <div
                className="flex cursor-pointer items-center gap-1 transition-colors hover:text-foreground"
                onClick={handleHeaderClick(header.column.id, header.column.getToggleSortingHandler())}
              >
                {flexRender(
                  header.column.columnDef.header,
                  header.getContext()
                )}
                {getIsSorted(header.column.id, header.column.getIsSorted()) === "asc" && (
                  <ArrowUp className="h-3 w-3 text-primary" />
                )}
                {getIsSorted(header.column.id, header.column.getIsSorted()) === "desc" && (
                  <ArrowDown className="h-3 w-3 text-primary" />
                )}
              </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 transition-colors hover:bg-primary/60 ${
                  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={`tusk-grid-row absolute left-0 flex transition-colors ${
                virtualRow.index % 2 === 1 ? "tusk-grid-row-odd" : ""
              }`}
              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-r border-border/15 text-xs"
                    style={{ width: w, minWidth: w }}
                  >
                    {flexRender(
                      cell.column.columnDef.cell,
                      cell.getContext()
                    )}
                  </div>
                );
              })}
            </div>
          );
        })}
      </div>
    </div>
  );
}
