tusk/src/components/table-viewer/TableDataView.tsx

import { useState, useCallback, useMemo } from "react";
import { useTableData } from "@/hooks/use-table-data";
import { ResultsTable } from "@/components/results/ResultsTable";
import { PaginationControls } from "./PaginationControls";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { updateRow as updateRowApi } from "@/lib/tauri";
import { getTableColumns } from "@/lib/tauri";
import { useQuery, useQueryClient } from "@tanstack/react-query";
import { useAppStore } from "@/stores/app-store";
import { toast } from "sonner";
import { Save, RotateCcw, Filter, Loader2, Lock } from "lucide-react";

interface Props {
  connectionId: string;
  schema: string;
  table: string;
}

export function TableDataView({ connectionId, schema, table }: Props) {
  const readOnlyMap = useAppStore((s) => s.readOnlyMap);
  const isReadOnly = readOnlyMap[connectionId] ?? true;

  const [page, setPage] = useState(1);
  const [pageSize, setPageSize] = useState(50);
  const [sortColumn, _setSortColumn] = useState<string | undefined>();
  const [sortDirection, _setSortDirection] = useState<string | undefined>();
  const [filter, setFilter] = useState("");
  const [appliedFilter, setAppliedFilter] = useState<string | undefined>();
  const [pendingChanges, setPendingChanges] = useState<
    Map<string, { rowIndex: number; colIndex: number; value: string | null }>
  >(new Map());
  const [isSaving, setIsSaving] = useState(false);
  const queryClient = useQueryClient();

  const { data, isLoading, error } = useTableData({
    connectionId,
    schema,
    table,
    page,
    pageSize,
    sortColumn,
    sortDirection,
    filter: appliedFilter,
  });

  const { data: columnsInfo } = useQuery({
    queryKey: ["table-columns", connectionId, schema, table],
    queryFn: () => getTableColumns(connectionId, schema, table),
  });

  const pkColumns = useMemo(
    () => columnsInfo?.filter((c) => c.is_primary_key).map((c) => c.name) ?? [],
    [columnsInfo]
  );

  const highlightedCells = useMemo(
    () => new Set(pendingChanges.keys()),
    [pendingChanges]
  );

  const handleCellDoubleClick = useCallback(
    (rowIndex: number, colIndex: number, value: unknown) => {
      if (isReadOnly) {
        toast.warning("Read-only mode is active", {
          description: "Switch to Read-Write mode to edit data.",
        });
        return;
      }
      const key = `${rowIndex}:${colIndex}`;
      const currentValue = pendingChanges.get(key)?.value ?? value;
      const newVal = prompt("Edit value (leave empty and click NULL for null):",
        currentValue === null ? "" : String(currentValue));
      if (newVal !== null) {
        setPendingChanges((prev) => {
          const next = new Map(prev);
          next.set(key, { rowIndex, colIndex, value: newVal === "" ? null : newVal });
          return next;
        });
      }
    },
    [pendingChanges, isReadOnly]
  );

  const handleCommit = async () => {
    if (!data || pkColumns.length === 0) {
      toast.error("Cannot save: no primary key detected");
      return;
    }
    setIsSaving(true);
    try {
      for (const [_key, change] of pendingChanges) {
        const row = data.rows[change.rowIndex];
        const pkValues = pkColumns.map((pkCol) => {
          const idx = data.columns.indexOf(pkCol);
          return row[idx];
        });
        const colName = data.columns[change.colIndex];

        await updateRowApi({
          connectionId,
          schema,
          table,
          pkColumns,
          pkValues: pkValues as unknown[],
          column: colName,
          value: change.value,
        });
      }
      setPendingChanges(new Map());
      queryClient.invalidateQueries({
        queryKey: ["table-data", connectionId],
      });
      toast.success(`${pendingChanges.size} change(s) saved`);
    } catch (err) {
      toast.error("Save failed", { description: String(err) });
    } finally {
      setIsSaving(false);
    }
  };

  const handleRollback = () => {
    setPendingChanges(new Map());
  };

  const handleApplyFilter = () => {
    setAppliedFilter(filter || undefined);
    setPage(1);
  };

  return (
    <div className="flex h-full flex-col">
      <div className="flex items-center gap-2 border-b px-2 py-1">
        {isReadOnly && (
          <span className="flex items-center gap-1 rounded bg-yellow-500/10 px-1.5 py-0.5 text-[10px] font-semibold text-yellow-600 dark:text-yellow-500">
            <Lock className="h-3 w-3" />
            Read-Only
          </span>
        )}
        <Filter className="h-3.5 w-3.5 text-muted-foreground" />
        <Input
          placeholder="WHERE clause (e.g. id > 10)"
          className="h-6 flex-1 text-xs"
          value={filter}
          onChange={(e) => setFilter(e.target.value)}
          onKeyDown={(e) => e.key === "Enter" && handleApplyFilter()}
        />
        <Button
          size="sm"
          variant="ghost"
          className="h-6 text-xs"
          onClick={handleApplyFilter}
        >
          Apply
        </Button>
        {pendingChanges.size > 0 && (
          <>
            <Button
              size="sm"
              variant="default"
              className="h-6 gap-1 text-xs"
              onClick={handleCommit}
              disabled={isSaving}
            >
              {isSaving ? (
                <Loader2 className="h-3 w-3 animate-spin" />
              ) : (
                <Save className="h-3 w-3" />
              )}
              Commit ({pendingChanges.size})
            </Button>
            <Button
              size="sm"
              variant="ghost"
              className="h-6 gap-1 text-xs"
              onClick={handleRollback}
            >
              <RotateCcw className="h-3 w-3" />
              Rollback
            </Button>
          </>
        )}
      </div>

      <div className="flex-1 overflow-hidden">
        {isLoading && !data ? (
          <div className="flex h-full items-center justify-center text-muted-foreground">
            <Loader2 className="mr-2 h-5 w-5 animate-spin" />
            Loading...
          </div>
        ) : error ? (
          <div className="p-4 text-sm text-destructive">
            {String(error)}
          </div>
        ) : data ? (
          <ResultsTable
            columns={data.columns}
            types={data.types}
            rows={data.rows}
            onCellDoubleClick={handleCellDoubleClick}
            highlightedCells={highlightedCells}
          />
        ) : null}
      </div>

      {data && (
        <PaginationControls
          page={data.page}
          pageSize={data.page_size}
          totalRows={data.total_rows}
          onPageChange={setPage}
          onPageSizeChange={(size) => {
            setPageSize(size);
            setPage(1);
          }}
        />
      )}
    </div>
  );
}