tusk/src/components/management/GrantRevokeDialog.tsx

import { useState } from "react";
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogFooter,
} from "@/components/ui/dialog";
import { Button } from "@/components/ui/button";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import {
  useRoles,
  useGrantRevoke,
  useTablePrivileges,
} from "@/hooks/use-management";
import { toast } from "sonner";
import { Loader2 } from "lucide-react";
import { Badge } from "@/components/ui/badge";

const PRIVILEGE_OPTIONS: Record<string, string[]> = {
  TABLE: ["SELECT", "INSERT", "UPDATE", "DELETE", "TRUNCATE", "REFERENCES", "TRIGGER", "ALL"],
  SCHEMA: ["USAGE", "CREATE"],
  DATABASE: ["CONNECT", "CREATE", "TEMPORARY"],
};

interface Props {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  connectionId: string;
  objectType: string;
  objectName: string;
  schema?: string;
  table?: string;
}

export function GrantRevokeDialog({
  open,
  onOpenChange,
  connectionId,
  objectType,
  objectName,
  schema,
  table,
}: Props) {
  const [action, setAction] = useState("GRANT");
  const [roleName, setRoleName] = useState("");
  const [privileges, setPrivileges] = useState<string[]>([]);
  const [withGrantOption, setWithGrantOption] = useState(false);

  const { data: roles } = useRoles(open ? connectionId : null);
  const { data: existingPrivileges } = useTablePrivileges(
    open && objectType === "TABLE" ? connectionId : null,
    schema ?? null,
    table ?? null
  );
  const grantRevokeMutation = useGrantRevoke();

  const [prevOpen, setPrevOpen] = useState(false);
  if (open !== prevOpen) {
    setPrevOpen(open);
    if (open) {
      setAction("GRANT");
      setRoleName("");
      setPrivileges([]);
      setWithGrantOption(false);
    }
  }

  const availablePrivileges = PRIVILEGE_OPTIONS[objectType.toUpperCase()] ?? PRIVILEGE_OPTIONS.TABLE;

  const togglePrivilege = (priv: string) => {
    setPrivileges((prev) =>
      prev.includes(priv)
        ? prev.filter((p) => p !== priv)
        : [...prev, priv]
    );
  };

  const handleSubmit = () => {
    if (!roleName) {
      toast.error("Please select a role");
      return;
    }
    if (privileges.length === 0) {
      toast.error("Please select at least one privilege");
      return;
    }
    grantRevokeMutation.mutate(
      {
        connectionId,
        params: {
          action,
          privileges,
          object_type: objectType,
          object_name: objectName,
          role_name: roleName,
          with_grant_option: withGrantOption,
        },
      },
      {
        onSuccess: () => {
          toast.success(
            `${action === "GRANT" ? "Granted" : "Revoked"} privileges on ${objectName}`
          );
          onOpenChange(false);
        },
        onError: (err) => {
          toast.error("Operation failed", { description: String(err) });
        },
      }
    );
  };

  return (
    <Dialog open={open} onOpenChange={onOpenChange}>
      <DialogContent className="sm:max-w-[520px]">
        <DialogHeader>
          <DialogTitle>Manage Privileges</DialogTitle>
        </DialogHeader>

        <div className="grid gap-3 py-4">
          <div className="grid grid-cols-4 items-center gap-3">
            <label className="text-right text-sm text-muted-foreground">Object</label>
            <div className="col-span-3 text-sm">
              <Badge variant="outline">{objectType}</Badge>{" "}
              <span className="font-medium">{objectName}</span>
            </div>
          </div>

          <div className="grid grid-cols-4 items-center gap-3">
            <label className="text-right text-sm text-muted-foreground">Action</label>
            <div className="col-span-3 flex gap-2">
              <Button
                size="sm"
                variant={action === "GRANT" ? "default" : "outline"}
                onClick={() => setAction("GRANT")}
              >
                Grant
              </Button>
              <Button
                size="sm"
                variant={action === "REVOKE" ? "default" : "outline"}
                onClick={() => setAction("REVOKE")}
              >
                Revoke
              </Button>
            </div>
          </div>

          <div className="grid grid-cols-4 items-center gap-3">
            <label className="text-right text-sm text-muted-foreground">Role</label>
            <Select value={roleName} onValueChange={setRoleName}>
              <SelectTrigger className="col-span-3">
                <SelectValue placeholder="Select role..." />
              </SelectTrigger>
              <SelectContent>
                {roles?.map((r) => (
                  <SelectItem key={r.name} value={r.name}>
                    {r.name}
                  </SelectItem>
                ))}
              </SelectContent>
            </Select>
          </div>

          <div className="grid grid-cols-4 items-start gap-3">
            <label className="text-right text-sm text-muted-foreground pt-1">Privileges</label>
            <div className="col-span-3 flex flex-wrap gap-1.5">
              {availablePrivileges.map((priv) => (
                <button
                  key={priv}
                  type="button"
                  className={`rounded-full border px-2.5 py-0.5 text-xs transition-colors ${
                    privileges.includes(priv)
                      ? "border-primary bg-primary text-primary-foreground"
                      : "border-border text-muted-foreground hover:text-foreground"
                  }`}
                  onClick={() => togglePrivilege(priv)}
                >
                  {priv}
                </button>
              ))}
            </div>
          </div>

          {action === "GRANT" && (
            <div className="grid grid-cols-4 items-center gap-3">
              <label className="text-right text-sm text-muted-foreground">Options</label>
              <label className="col-span-3 flex items-center gap-2 text-sm cursor-pointer">
                <input
                  type="checkbox"
                  checked={withGrantOption}
                  onChange={(e) => setWithGrantOption(e.target.checked)}
                  className="rounded border-input"
                />
                WITH GRANT OPTION
              </label>
            </div>
          )}

          {existingPrivileges && existingPrivileges.length > 0 && (
            <div className="grid grid-cols-4 items-start gap-3">
              <label className="text-right text-sm text-muted-foreground pt-1">Current</label>
              <div className="col-span-3 max-h-32 overflow-y-auto rounded border p-2">
                <div className="space-y-1">
                  {existingPrivileges.map((p, i) => (
                    <div key={i} className="flex items-center gap-2 text-xs">
                      <span className="font-medium">{p.grantee}</span>
                      <Badge variant="secondary" className="text-[10px]">
                        {p.privilege_type}
                      </Badge>
                      {p.is_grantable && (
                        <Badge variant="outline" className="text-[10px]">
                          GRANTABLE
                        </Badge>
                      )}
                    </div>
                  ))}
                </div>
              </div>
            </div>
          )}
        </div>

        <DialogFooter>
          <Button variant="outline" onClick={() => onOpenChange(false)}>
            Cancel
          </Button>
          <Button onClick={handleSubmit} disabled={grantRevokeMutation.isPending}>
            {grantRevokeMutation.isPending && (
              <Loader2 className="mr-2 h-4 w-4 animate-spin" />
            )}
            {action === "GRANT" ? "Grant" : "Revoke"}
          </Button>
        </DialogFooter>
      </DialogContent>
    </Dialog>
  );
}