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 { ResultsJsonView } from "@/components/results/ResultsJsonView";
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, Download, Plus, Table2, Braces } from "lucide-react";
import { InsertRowDialog } from "./InsertRowDialog";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { exportCsv, exportJson } from "@/lib/tauri";
import { save } from "@tauri-apps/plugin-dialog";

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 [insertDialogOpen, setInsertDialogOpen] = useState(false);
  const [viewMode, setViewMode] = useState<"table" | "json">("table");
  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 usesCtid = pkColumns.length === 0;

  const handleCommit = async () => {
    if (!data) return;
    if (pkColumns.length === 0 && (!data.ctids || data.ctids.length === 0)) {
      toast.error("Cannot save: no primary key and no ctid available");
      return;
    }
    setIsSaving(true);
    try {
      for (const [, change] of pendingChanges) {
        const row = data.rows[change.rowIndex];
        const colName = data.columns[change.colIndex];

        if (usesCtid) {
          await updateRowApi({
            connectionId,
            schema,
            table,
            pkColumns: [],
            pkValues: [],
            column: colName,
            value: change.value,
            ctid: data.ctids[change.rowIndex],
          });
        } else {
          const pkValues = pkColumns.map((pkCol) => {
            const idx = data.columns.indexOf(pkCol);
            return row[idx];
          });
          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 handleExport = useCallback(
    async (format: "csv" | "json") => {
      if (!data || data.columns.length === 0) return;
      const ext = format === "csv" ? "csv" : "json";
      const path = await save({
        title: `Export as ${ext.toUpperCase()}`,
        defaultPath: `${table}.${ext}`,
        filters: [{ name: ext.toUpperCase(), extensions: [ext] }],
      });
      if (!path) return;
      try {
        if (format === "csv") {
          await exportCsv(path, data.columns, data.rows as unknown[][]);
        } else {
          await exportJson(path, data.columns, data.rows as unknown[][]);
        }
        toast.success(`Exported to ${path}`);
      } catch (err) {
        toast.error("Export failed", { description: String(err) });
      }
    },
    [data, table]
  );

  const handleSort = useCallback(
    (column: string | undefined, direction: string | undefined) => {
      setSortColumn(column);
      setSortDirection(direction);
      setPage(1);
    },
    []
  );

  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>
        )}
        {!isReadOnly && usesCtid && (
          <span
            className="rounded bg-orange-500/10 px-1.5 py-0.5 text-[10px] font-medium text-orange-600 dark:text-orange-400"
            title="This table has no primary key. Edits use physical row ID (ctid), which may change after VACUUM or concurrent writes."
          >
            No PK — using ctid
          </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>
        {data && data.columns.length > 0 && (
          <>
            <DropdownMenu>
              <DropdownMenuTrigger asChild>
                <Button
                  size="sm"
                  variant="ghost"
                  className="h-6 gap-1 text-xs"
                >
                  <Download className="h-3 w-3" />
                  Export
                </Button>
              </DropdownMenuTrigger>
              <DropdownMenuContent align="start">
                <DropdownMenuItem onClick={() => handleExport("csv")}>
                  Export CSV
                </DropdownMenuItem>
                <DropdownMenuItem onClick={() => handleExport("json")}>
                  Export JSON
                </DropdownMenuItem>
              </DropdownMenuContent>
            </DropdownMenu>
            <div className="flex items-center rounded-md border text-xs">
              <button
                className={`flex items-center gap-1 px-2 py-0.5 font-medium ${
                  viewMode === "table"
                    ? "bg-muted text-foreground"
                    : "text-muted-foreground hover:text-foreground"
                }`}
                onClick={() => setViewMode("table")}
                title="Table view"
              >
                <Table2 className="h-3 w-3" />
                Table
              </button>
              <button
                className={`flex items-center gap-1 px-2 py-0.5 font-medium ${
                  viewMode === "json"
                    ? "bg-muted text-foreground"
                    : "text-muted-foreground hover:text-foreground"
                }`}
                onClick={() => setViewMode("json")}
                title="JSON view"
              >
                <Braces className="h-3 w-3" />
                JSON
              </button>
            </div>
          </>
        )}
        {!isReadOnly && (
          <Button
            size="sm"
            variant="ghost"
            className="h-6 gap-1 text-xs"
            onClick={() => setInsertDialogOpen(true)}
          >
            <Plus className="h-3 w-3" />
            Insert Row
          </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 ? (
          viewMode === "json" ? (
            <ResultsJsonView
              columns={data.columns}
              rows={data.rows}
            />
          ) : (
            <ResultsTable
              columns={data.columns}
              types={data.types}
              rows={data.rows}
              onCellDoubleClick={handleCellDoubleClick}
              highlightedCells={highlightedCells}
              externalSort={{
                column: sortColumn,
                direction: sortDirection,
                onSort: handleSort,
              }}
            />
          )
        ) : null}
      </div>

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

      <InsertRowDialog
        open={insertDialogOpen}
        onOpenChange={setInsertDialogOpen}
        connectionId={connectionId}
        schema={schema}
        table={table}
        onSuccess={() => {
          queryClient.invalidateQueries({
            queryKey: ["table-data", connectionId],
          });
        }}
      />
    </div>
  );
}