"use client";

import {
  type archestraApiTypes,
  E2eTestId,
  getDeleteVirtualKeyButtonTestId,
  getVirtualKeyRowTestId,
} from "@archestra/shared";
import type { ColumnDef } from "@tanstack/react-table";
import {
  Globe,
  Key,
  Loader2,
  Pencil,
  Plus,
  Trash2,
  User,
  Users,
} from "lucide-react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { CopyableCode } from "@/components/copyable-code";
import { DeleteConfirmDialog } from "@/components/delete-confirm-dialog";
import { ExpirationDateTimeField } from "@/components/expiration-date-time-field";
import { FormDialog } from "@/components/form-dialog";
import { LlmProviderApiKeyDropdown } from "@/components/llm-provider-api-key-dropdown";
import type { LlmProviderApiKeyResponse } from "@/components/llm-provider-api-key-form";
import {
  formatProviderKeySummary,
  type ProviderApiKeyMap,
  providerApiKeyMapToArray,
} from "@/components/provider-key-mappings-field";
import { ProviderKeyAccessFields } from "@/components/proxy-auth-provider-key-fields";
import { QueryLoadError } from "@/components/query-load-error";
import { ResourceVisibilityBadge } from "@/components/resource-visibility-badge";
import { SearchInput } from "@/components/search-input";
import { TableRowActions } from "@/components/table-row-actions";
import { Badge } from "@/components/ui/badge";
import { Button } from "@/components/ui/button";
import { DataTable } from "@/components/ui/data-table";
import {
  DialogBody,
  DialogForm,
  DialogStickyFooter,
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { MultiSelectCombobox } from "@/components/ui/multi-select-combobox";
import { RadioGroup, RadioGroupItem } from "@/components/ui/radio-group";
import {
  Select,
  SelectContent,
  SelectItem,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import { DialogCancelButton } from "@/components/unsaved-changes-guard";
import { hasUnsavedChanges } from "@/components/unsaved-changes-guard-utils";
import {
  type VisibilityOption,
  VisibilitySelector,
} from "@/components/visibility-selector";
import { useHasPermissions, useSession } from "@/lib/auth/auth.query";
import { useFeature } from "@/lib/config/config.query";
import { useDataTableQueryParams } from "@/lib/hooks/use-data-table-query-params";
import { useLlmProviderApiKeys } from "@/lib/llm-provider-api-keys.query";
import { useTeams } from "@/lib/teams/team.query";
import { formatRelativeTime } from "@/lib/utils/date-time";
import {
  useAllVirtualApiKeys,
  useCreateVirtualApiKey,
  useDeleteVirtualApiKey,
  useUpdateVirtualApiKey,
} from "@/lib/virtual-api-keys.query";
import { useSetCredentialsAction } from "../layout";
import { OwnerSelectField, shouldShowOwnerField } from "./owner-select-field";

type VirtualKeyWithParent =
  archestraApiTypes.GetAllVirtualApiKeysResponses["200"]["data"][number];
type VirtualKeyScope = NonNullable<
  archestraApiTypes.CreateVirtualApiKeyData["body"]["scope"]
>;
type VirtualKeyType = NonNullable<
  archestraApiTypes.CreateVirtualApiKeyData["body"]["keyType"]
>;

const KEY_TYPE_OPTIONS: {
  value: VirtualKeyType;
  label: string;
  description: string;
}[] = [
  {
    value: "standard",
    label: "Standard",
    description:
      "Maps to provider API keys and is sent in the Authorization header as a provider key replacement.",
  },
  {
    value: "passthrough",
    label: "Passthrough",
    description:
      "Carries no provider key. Sent in the X-Archestra-Virtual-Key header to attribute a request to a user when the provider credential is passed through (e.g. a Claude Code subscription token).",
  },
];

const KEY_TYPE_LABEL: Record<VirtualKeyType, string> = {
  standard: "Standard",
  passthrough: "Passthrough",
};

export default function VirtualKeysPage() {
  const {
    searchParams,
    pageIndex,
    pageSize,
    offset,
    setPagination,
    updateQueryParams,
  } = useDataTableQueryParams();
  const search = searchParams.get("search") || "";
  const providerApiKeyIdFilter = searchParams.get("providerApiKeyId") || "all";
  const keyTypeFilter = searchParams.get("keyType") || "all";

  const {
    data: response,
    isPending,
    isLoadingError: isVirtualKeysLoadError,
    refetch: refetchVirtualKeys,
  } = useAllVirtualApiKeys({
    limit: pageSize,
    offset,
    search: search || undefined,
    providerApiKeyId:
      providerApiKeyIdFilter === "all" ? undefined : providerApiKeyIdFilter,
    keyType:
      keyTypeFilter === "all" ? undefined : (keyTypeFilter as VirtualKeyType),
    toastOnError: false,
  });
  const virtualKeys = response?.data ?? [];
  const paginationMeta = response?.pagination;

  const { data: apiKeys = [] } = useLlmProviderApiKeys();
  const { data: session } = useSession();
  const { data: canReadTeams } = useHasPermissions({ team: ["read"] });
  const { data: isVirtualKeyAdmin } = useHasPermissions({
    llmVirtualKey: ["admin"],
  });
  const { data: teams = [] } = useTeams({ enabled: !!canReadTeams });
  const defaultExpirationSeconds = useFeature(
    "virtualKeyDefaultExpirationSeconds",
  );

  const [isCreateDialogOpen, setIsCreateDialogOpen] = useState(false);
  const [editingKey, setEditingKey] = useState<VirtualKeyWithParent | null>(
    null,
  );
  const [providerApiKeyFilterOpen, setProviderApiKeyFilterOpen] =
    useState(false);
  const [isDeleteDialogOpen, setIsDeleteDialogOpen] = useState(false);
  const [deletingKey, setDeletingKey] = useState<VirtualKeyWithParent | null>(
    null,
  );

  const columns: ColumnDef<VirtualKeyWithParent>[] = useMemo(
    () => [
      {
        accessorKey: "name",
        header: "Name",
        cell: ({ row }) => (
          <span
            className="font-medium"
            data-testid={getVirtualKeyRowTestId(row.original.name)}
          >
            {row.original.name}
          </span>
        ),
      },
      {
        accessorKey: "tokenStart",
        header: "Token",
        cell: ({ row }) => (
          <code className="text-xs text-muted-foreground">
            {row.original.tokenStart}...
          </code>
        ),
      },
      {
        id: "keyType",
        header: "Type",
        cell: ({ row }) => (
          <Badge
            variant={
              row.original.keyType === "passthrough" ? "outline" : "secondary"
            }
          >
            {KEY_TYPE_LABEL[row.original.keyType]}
          </Badge>
        ),
      },
      {
        id: "accessibleTo",
        header: "Accessible to",
        cell: ({ row }) => (
          <ResourceVisibilityBadge
            scope={row.original.scope as VirtualKeyScope | undefined}
            teams={row.original.teams}
            authorId={row.original.authorId}
            authorName={row.original.authorName}
            currentUserId={session?.user?.id}
            // The "Accessible to" column also lists team- and org-scoped keys, so
            // label the current user's own personal key "Me" (rather than leaving
            // it blank) to keep every row consistently attributed.
            showSelfAsMe
          />
        ),
      },
      {
        id: "providerKeys",
        header: "Provider Keys",
        cell: ({ row }) => (
          <span className="text-sm text-muted-foreground">
            {row.original.keyType === "passthrough"
              ? "N/A"
              : formatProviderKeySummary(row.original.providerApiKeys)}
          </span>
        ),
      },
      {
        accessorKey: "expiresAt",
        header: "Expires",
        cell: ({ row }) => (
          <span className="text-sm text-muted-foreground">
            {formatExpiration(row.original.expiresAt)}
          </span>
        ),
      },
      {
        accessorKey: "lastUsedAt",
        header: "Last Used",
        cell: ({ row }) =>
          row.original.lastUsedAt ? (
            <span className="text-sm text-muted-foreground">
              {new Date(row.original.lastUsedAt).toLocaleDateString()}
            </span>
          ) : (
            <span className="text-sm text-muted-foreground">Never</span>
          ),
      },
      {
        id: "actions",
        header: "Actions",
        cell: ({ row }) => (
          <TableRowActions
            actions={[
              {
                icon: <Pencil className="h-4 w-4" />,
                label: "Edit",
                onClick: () => setEditingKey(row.original),
              },
              {
                icon: <Trash2 className="h-4 w-4" />,
                label: "Delete",
                variant: "destructive",
                testId: getDeleteVirtualKeyButtonTestId(row.original.name),
                onClick: () => {
                  setDeletingKey(row.original);
                  setIsDeleteDialogOpen(true);
                },
              },
            ]}
          />
        ),
      },
    ],
    [session?.user?.id],
  );

  const parentableKeys = apiKeys;

  const visibilityOptions = useMemo(
    () =>
      getVirtualKeyVisibilityOptions({
        canReadTeams: !!canReadTeams,
        isAdmin: !!isVirtualKeyAdmin,
      }),
    [canReadTeams, isVirtualKeyAdmin],
  );

  const setCredentialsAction = useSetCredentialsAction();
  useEffect(() => {
    setCredentialsAction(
      <Button
        onClick={() => setIsCreateDialogOpen(true)}
        data-testid={E2eTestId.AddVirtualKeyButton}
      >
        <Plus className="h-4 w-4" />
        Create Virtual Key
      </Button>,
    );
    return () => setCredentialsAction(null);
  }, [setCredentialsAction]);

  return (
    <>
      <div
        className="mb-4 flex flex-wrap gap-4"
        data-testid={E2eTestId.VirtualKeysPage}
      >
        <SearchInput
          objectNamePlural="virtual keys"
          searchFields={["name"]}
          paramName="search"
        />
        <LlmProviderApiKeyDropdown
          availableKeys={parentableKeys}
          selectedApiKeyId={
            providerApiKeyIdFilter === "all" ? null : providerApiKeyIdFilter
          }
          open={providerApiKeyFilterOpen}
          onOpenChange={setProviderApiKeyFilterOpen}
          onSelectKey={(value) => {
            updateQueryParams({
              providerApiKeyId: value,
              page: "1",
            });
            setProviderApiKeyFilterOpen(false);
          }}
          triggerVariant="select"
          triggerClassName="w-full sm:w-[280px] h-9 text-sm"
          popoverClassName="w-[var(--radix-popover-trigger-width)]"
          allOptionLabel="All provider API keys"
          allOptionSelected={providerApiKeyIdFilter === "all"}
          onSelectAllOption={() => {
            updateQueryParams({
              providerApiKeyId: null,
              page: "1",
            });
            setProviderApiKeyFilterOpen(false);
          }}
        />
        <Select
          value={keyTypeFilter}
          onValueChange={(value) =>
            updateQueryParams({
              keyType: value === "all" ? null : value,
              page: "1",
            })
          }
        >
          <SelectTrigger className="h-9 w-full text-sm sm:w-[200px]">
            <SelectValue placeholder="All key types" />
          </SelectTrigger>
          <SelectContent>
            <SelectItem value="all">All key types</SelectItem>
            <SelectItem value="standard">Standard</SelectItem>
            <SelectItem value="passthrough">Passthrough</SelectItem>
          </SelectContent>
        </Select>
      </div>

      {isVirtualKeysLoadError ? (
        <QueryLoadError
          title="Couldn't load your virtual keys"
          onRetry={() => refetchVirtualKeys()}
        />
      ) : (
        <DataTable
          columns={columns}
          data={virtualKeys}
          getRowId={(row) => row.id}
          hideSelectedCount
          isLoading={isPending}
          emptyMessage={
            parentableKeys.length === 0
              ? "Add an API key first to create virtual keys"
              : "No virtual keys yet"
          }
          manualPagination
          pagination={{
            pageIndex,
            pageSize,
            total: paginationMeta?.total ?? 0,
          }}
          onPaginationChange={setPagination}
          hasActiveFilters={Boolean(
            search ||
              providerApiKeyIdFilter !== "all" ||
              keyTypeFilter !== "all",
          )}
          filteredEmptyMessage="No virtual keys match your filters. Try adjusting your search."
          onClearFilters={() =>
            updateQueryParams({
              search: null,
              providerApiKeyId: null,
              keyType: null,
              page: "1",
            })
          }
        />
      )}

      <CreateVirtualKeyDialog
        open={isCreateDialogOpen}
        onOpenChange={setIsCreateDialogOpen}
        parentableKeys={parentableKeys}
        defaultExpirationSeconds={defaultExpirationSeconds ?? null}
        visibilityOptions={visibilityOptions}
        teams={teams}
        canReadTeams={!!canReadTeams}
        isVirtualKeyAdmin={!!isVirtualKeyAdmin}
      />

      <EditVirtualKeyDialog
        open={!!editingKey}
        onOpenChange={(open) => !open && setEditingKey(null)}
        virtualKey={editingKey}
        providerApiKeys={parentableKeys}
        visibilityOptions={visibilityOptions}
        teams={teams}
        canReadTeams={!!canReadTeams}
      />

      <DeleteVirtualKeyDialog
        open={isDeleteDialogOpen}
        onOpenChange={setIsDeleteDialogOpen}
        virtualKey={deletingKey}
      />
    </>
  );
}

function CreateVirtualKeyDialog({
  open,
  onOpenChange,
  parentableKeys,
  defaultExpirationSeconds,
  visibilityOptions,
  teams,
  canReadTeams,
  isVirtualKeyAdmin,
}: {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  parentableKeys: LlmProviderApiKeyResponse[];
  defaultExpirationSeconds: number | null;
  visibilityOptions: VisibilityOption<VirtualKeyScope>[];
  teams: Array<{ id: string; name: string }>;
  canReadTeams: boolean;
  isVirtualKeyAdmin: boolean;
}) {
  const createMutation = useCreateVirtualApiKey();

  const [keyType, setKeyType] = useState<VirtualKeyType>("standard");
  const [newKeyName, setNewKeyName] = useState("");
  const [ownerId, setOwnerId] = useState("");
  const [expiresAt, setExpiresAt] = useState<Date | null>(null);
  const [scope, setScope] = useState<VirtualKeyScope>(
    getDefaultVirtualKeyScope(visibilityOptions),
  );
  const [teamIds, setTeamIds] = useState<string[]>([]);
  const [providerApiKeyIds, setProviderApiKeyIds] = useState<ProviderApiKeyMap>(
    {},
  );
  const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
  const [createdKeyExpiresAt, setCreatedKeyExpiresAt] = useState<Date | null>(
    null,
  );

  const prevOpenRef = useRef(open);
  const initialSnapshotRef = useRef<Record<string, unknown> | null>(null);

  useEffect(() => {
    const wasOpen = prevOpenRef.current;
    prevOpenRef.current = open;
    if (open && !wasOpen) {
      setCreatedKeyValue(null);
      setCreatedKeyExpiresAt(null);
      const initialExpiresAt = computeDefaultExpiresAt(
        defaultExpirationSeconds,
      );
      const initialScope = getDefaultVirtualKeyScope(visibilityOptions);
      setKeyType("standard");
      setNewKeyName("");
      setExpiresAt(initialExpiresAt);
      setScope(initialScope);
      setTeamIds([]);
      setProviderApiKeyIds({});
      setOwnerId("");
      initialSnapshotRef.current = {
        keyType: "standard",
        newKeyName: "",
        ownerId: "",
        expiresAt: initialExpiresAt,
        scope: initialScope,
        teamIds: [],
        providerApiKeyIds: {},
      };
    }
  }, [open, defaultExpirationSeconds, visibilityOptions]);

  const isPassthrough = keyType === "passthrough";
  // Passthrough keys are always personal. Admins can mint a key on behalf of
  // another org member; left unset, the key belongs to the creator.
  const showOwnerField = shouldShowOwnerField(
    isVirtualKeyAdmin,
    isPassthrough ? "personal" : scope,
  );
  const standardReady =
    (scope !== "team" || teamIds.length > 0) &&
    providerApiKeyMapToArray(providerApiKeyIds).length > 0;
  const canSubmit =
    newKeyName.trim().length > 0 &&
    (isPassthrough || standardReady) &&
    !createMutation.isPending;

  // Once the key is created the form is replaced by the reveal view, so there
  // is nothing left to lose — only guard the editable form.
  const isDirty =
    !createdKeyValue &&
    initialSnapshotRef.current !== null &&
    hasUnsavedChanges(initialSnapshotRef.current, {
      keyType,
      newKeyName,
      ownerId,
      expiresAt,
      scope,
      teamIds: [...teamIds].sort(),
      providerApiKeyIds,
    });

  const handleCreate = useCallback(async () => {
    if (!newKeyName.trim()) return;
    const owner = showOwnerField && ownerId ? ownerId : undefined;
    try {
      const result = await createMutation.mutateAsync({
        data: isPassthrough
          ? {
              name: newKeyName.trim(),
              keyType: "passthrough",
              expiresAt: expiresAt ?? undefined,
              ownerId: owner,
            }
          : {
              name: newKeyName.trim(),
              keyType: "standard",
              expiresAt: expiresAt ?? undefined,
              scope,
              teams: scope === "team" ? teamIds : [],
              providerApiKeys: providerApiKeyMapToArray(providerApiKeyIds),
              ownerId: owner,
            },
      });
      setNewKeyName("");
      if (result?.value) {
        setCreatedKeyValue(result.value);
        setCreatedKeyExpiresAt(expiresAt);
      }
    } catch {
      // handled by mutation
    }
  }, [
    createMutation,
    expiresAt,
    isPassthrough,
    providerApiKeyIds,
    newKeyName,
    scope,
    teamIds,
    showOwnerField,
    ownerId,
  ]);

  return (
    <FormDialog
      open={open}
      onOpenChange={onOpenChange}
      title={
        createdKeyValue ? "Virtual API Key Created" : "Create Virtual API Key"
      }
      description={
        createdKeyValue
          ? undefined
          : "Map provider API keys, or create a passthrough key to attribute requests to a user."
      }
      size="medium"
      isDirty={isDirty}
    >
      <DialogForm onSubmit={handleCreate}>
        <DialogBody
          className="space-y-4"
          data-testid={E2eTestId.VirtualKeyCreateDialog}
        >
          {createdKeyValue ? (
            <div className="space-y-3">
              <div className="flex items-center gap-2 text-sm font-medium">
                <Key className="h-4 w-4" />
                Copy this key now. It won&apos;t be shown again.
              </div>
              <div data-testid={E2eTestId.VirtualKeyValue}>
                <CopyableCode value={createdKeyValue} />
              </div>
              <div className="text-sm text-muted-foreground">
                <span className="font-medium text-foreground">Expires:</span>{" "}
                {formatExpiration(createdKeyExpiresAt)}
              </div>
            </div>
          ) : (
            <>
              <div className="space-y-2">
                <Label htmlFor="virtual-key-name">Name</Label>
                <Input
                  id="virtual-key-name"
                  value={newKeyName}
                  onChange={(e) => setNewKeyName(e.target.value)}
                  placeholder="My virtual key"
                />
              </div>

              <KeyTypeField value={keyType} onChange={setKeyType} />

              {isPassthrough ? (
                <>
                  {showOwnerField && (
                    <OwnerSelectField value={ownerId} onChange={setOwnerId} />
                  )}

                  <div className="space-y-2">
                    <ExpirationDateTimeField
                      value={expiresAt}
                      onChange={setExpiresAt}
                      noExpirationText="Key will never expire"
                      formatExpiration={formatExpiration}
                    />
                  </div>
                </>
              ) : (
                <>
                  <VirtualKeyVisibilityField
                    value={scope}
                    onValueChange={(nextScope) => {
                      setScope(nextScope);
                      if (nextScope !== "team") {
                        setTeamIds([]);
                      }
                    }}
                    teamIds={teamIds}
                    onTeamIdsChange={setTeamIds}
                    teams={teams}
                    canReadTeams={canReadTeams}
                    visibilityOptions={visibilityOptions}
                  />

                  {showOwnerField && (
                    <OwnerSelectField value={ownerId} onChange={setOwnerId} />
                  )}

                  <div className="space-y-2">
                    <ExpirationDateTimeField
                      value={expiresAt}
                      onChange={setExpiresAt}
                      noExpirationText="Key will never expire"
                      formatExpiration={formatExpiration}
                    />
                  </div>

                  <ProviderKeyAccessFields
                    providerApiKeyIds={providerApiKeyIds}
                    onProviderApiKeyIdsChange={setProviderApiKeyIds}
                    providerApiKeys={parentableKeys}
                  />
                </>
              )}
            </>
          )}
        </DialogBody>
        <DialogStickyFooter className="mt-0">
          <DialogCancelButton>
            {createdKeyValue ? "Close" : "Cancel"}
          </DialogCancelButton>
          {!createdKeyValue && (
            <Button type="submit" disabled={!canSubmit}>
              {createMutation.isPending && (
                <Loader2 className="h-4 w-4 animate-spin" />
              )}
              Create
            </Button>
          )}
        </DialogStickyFooter>
      </DialogForm>
    </FormDialog>
  );
}

function EditVirtualKeyDialog({
  open,
  onOpenChange,
  virtualKey,
  providerApiKeys,
  visibilityOptions,
  teams,
  canReadTeams,
}: {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  virtualKey: VirtualKeyWithParent | null;
  providerApiKeys: LlmProviderApiKeyResponse[];
  visibilityOptions: VisibilityOption<VirtualKeyScope>[];
  teams: Array<{ id: string; name: string }>;
  canReadTeams: boolean;
}) {
  const updateMutation = useUpdateVirtualApiKey();
  const [name, setName] = useState("");
  const [expiresAt, setExpiresAt] = useState<Date | null>(null);
  const [scope, setScope] = useState<VirtualKeyScope>(
    getDefaultVirtualKeyScope(visibilityOptions),
  );
  const [teamIds, setTeamIds] = useState<string[]>([]);
  const [providerApiKeyIds, setProviderApiKeyIds] = useState<ProviderApiKeyMap>(
    {},
  );
  const initialSnapshotRef = useRef<Record<string, unknown> | null>(null);

  useEffect(() => {
    if (!open || !virtualKey) {
      return;
    }

    const initialName = virtualKey.name;
    const initialExpiresAt = virtualKey.expiresAt
      ? new Date(virtualKey.expiresAt)
      : null;
    const initialScope = (virtualKey.scope as VirtualKeyScope) ?? "personal";
    const initialTeamIds = virtualKey.teams.map((team) => team.id);
    const initialProviderApiKeyIds = Object.fromEntries(
      virtualKey.providerApiKeys.map((mapping) => [
        mapping.provider,
        mapping.providerApiKeyId,
      ]),
    );

    setName(initialName);
    setExpiresAt(initialExpiresAt);
    setScope(initialScope);
    setTeamIds(initialTeamIds);
    setProviderApiKeyIds(initialProviderApiKeyIds);
    initialSnapshotRef.current = {
      name: initialName,
      expiresAt: initialExpiresAt,
      scope: initialScope,
      teamIds: [...initialTeamIds].sort(),
      providerApiKeyIds: initialProviderApiKeyIds,
    };
  }, [open, virtualKey]);

  // The key type is fixed at creation; only its own configuration is editable.
  const isPassthrough = virtualKey?.keyType === "passthrough";

  const handleUpdate = useCallback(async () => {
    if (!virtualKey || !name.trim()) {
      return;
    }

    try {
      const result = await updateMutation.mutateAsync({
        id: virtualKey.id,
        data: isPassthrough
          ? {
              name: name.trim(),
              keyType: "passthrough",
              expiresAt: expiresAt ?? undefined,
            }
          : {
              name: name.trim(),
              keyType: "standard",
              expiresAt: expiresAt ?? undefined,
              scope,
              teams: scope === "team" ? teamIds : [],
              providerApiKeys: providerApiKeyMapToArray(providerApiKeyIds),
            },
      });

      if (result) {
        onOpenChange(false);
      }
    } catch {
      // handled by mutation
    }
  }, [
    expiresAt,
    isPassthrough,
    providerApiKeyIds,
    name,
    onOpenChange,
    scope,
    teamIds,
    updateMutation,
    virtualKey,
  ]);

  if (!virtualKey) {
    return null;
  }

  const standardReady =
    (scope !== "team" || teamIds.length > 0) &&
    providerApiKeyMapToArray(providerApiKeyIds).length > 0;
  const canSubmit =
    name.trim().length > 0 &&
    (isPassthrough || standardReady) &&
    !updateMutation.isPending;
  const isDirty =
    initialSnapshotRef.current !== null &&
    hasUnsavedChanges(initialSnapshotRef.current, {
      name,
      expiresAt,
      scope,
      teamIds: [...teamIds].sort(),
      providerApiKeyIds,
    });

  return (
    <FormDialog
      open={open}
      onOpenChange={onOpenChange}
      title="Edit Virtual API Key"
      description={
        isPassthrough
          ? "Update the passthrough key name and expiration."
          : "Update the virtual key name, visibility, and expiration."
      }
      size="medium"
      isDirty={isDirty}
    >
      <DialogForm onSubmit={handleUpdate}>
        <DialogBody className="space-y-4">
          <div className="space-y-2">
            <Label htmlFor="edit-virtual-key-name">Name</Label>
            <Input
              id="edit-virtual-key-name"
              value={name}
              onChange={(event) => setName(event.target.value)}
              placeholder="My virtual key"
            />
          </div>

          {isPassthrough ? (
            <div className="space-y-2">
              <ExpirationDateTimeField
                value={expiresAt}
                onChange={setExpiresAt}
                noExpirationText="Key will never expire"
                formatExpiration={formatExpiration}
              />
            </div>
          ) : (
            <>
              <VirtualKeyVisibilityField
                value={scope}
                onValueChange={(nextScope) => {
                  setScope(nextScope);
                  if (nextScope !== "team") {
                    setTeamIds([]);
                  }
                }}
                teamIds={teamIds}
                onTeamIdsChange={setTeamIds}
                teams={teams}
                canReadTeams={canReadTeams}
                visibilityOptions={visibilityOptions}
              />

              <div className="space-y-2">
                <ExpirationDateTimeField
                  value={expiresAt}
                  onChange={setExpiresAt}
                  noExpirationText="Key will never expire"
                  formatExpiration={formatExpiration}
                />
              </div>

              <ProviderKeyAccessFields
                providerApiKeyIds={providerApiKeyIds}
                onProviderApiKeyIdsChange={setProviderApiKeyIds}
                providerApiKeys={providerApiKeys}
              />
            </>
          )}
        </DialogBody>
        <DialogStickyFooter className="mt-0">
          <DialogCancelButton>Cancel</DialogCancelButton>
          <Button type="submit" disabled={!canSubmit}>
            {updateMutation.isPending && (
              <Loader2 className="h-4 w-4 animate-spin" />
            )}
            Save Changes
          </Button>
        </DialogStickyFooter>
      </DialogForm>
    </FormDialog>
  );
}

function DeleteVirtualKeyDialog({
  open,
  onOpenChange,
  virtualKey,
}: {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  virtualKey: VirtualKeyWithParent | null;
}) {
  const deleteMutation = useDeleteVirtualApiKey();

  return (
    <DeleteConfirmDialog
      open={open}
      onOpenChange={onOpenChange}
      title="Delete Virtual Key"
      description={
        <div data-testid={E2eTestId.VirtualKeyDeleteDialog}>
          Are you sure you want to delete "{virtualKey?.name}"? This action
          cannot be undone.
        </div>
      }
      confirmLabel="Delete"
      isPending={deleteMutation.isPending}
      onConfirm={() => {
        if (!virtualKey) return;

        deleteMutation.mutate(
          {
            id: virtualKey.id,
          },
          {
            onSuccess: () => {
              onOpenChange(false);
            },
          },
        );
      }}
    />
  );
}

function VirtualKeyVisibilityField({
  value,
  onValueChange,
  teamIds,
  onTeamIdsChange,
  teams,
  canReadTeams,
  visibilityOptions,
}: {
  value: VirtualKeyScope;
  onValueChange: (value: VirtualKeyScope) => void;
  teamIds: string[];
  onTeamIdsChange: (value: string[]) => void;
  teams: Array<{ id: string; name: string }>;
  canReadTeams: boolean;
  visibilityOptions: VisibilityOption<VirtualKeyScope>[];
}) {
  return (
    <VisibilitySelector
      heading="Who can use this virtual key"
      value={value}
      options={visibilityOptions}
      onValueChange={onValueChange}
    >
      {value === "team" && (
        <div className="space-y-2">
          <Label>Teams</Label>
          <MultiSelectCombobox
            disabled={!canReadTeams}
            options={teams.map((team) => ({
              value: team.id,
              label: team.name,
            }))}
            value={teamIds}
            onChange={onTeamIdsChange}
            placeholder={canReadTeams ? "Search teams..." : "Teams unavailable"}
            emptyMessage="No teams found."
          />
        </div>
      )}
    </VisibilitySelector>
  );
}

function KeyTypeField({
  value,
  onChange,
}: {
  value: VirtualKeyType;
  onChange: (value: VirtualKeyType) => void;
}) {
  return (
    <div className="space-y-2">
      <Label>Key type</Label>
      <RadioGroup
        value={value}
        onValueChange={(next) => onChange(next as VirtualKeyType)}
        className="gap-2"
      >
        {KEY_TYPE_OPTIONS.map((option) => (
          <Label
            key={option.value}
            htmlFor={`virtual-key-type-${option.value}`}
            className="flex cursor-pointer items-start gap-3 rounded-md border p-3 font-normal has-[:checked]:border-primary"
          >
            <RadioGroupItem
              id={`virtual-key-type-${option.value}`}
              value={option.value}
              className="mt-0.5"
            />
            <div className="space-y-1">
              <div className="font-medium">{option.label}</div>
              <p className="text-sm text-muted-foreground">
                {option.description}
              </p>
            </div>
          </Label>
        ))}
      </RadioGroup>
    </div>
  );
}

function formatExpiration(date: Date | string | null): string {
  return formatRelativeTime(date);
}

function computeDefaultExpiresAt(defaultSeconds: number | null): Date | null {
  if (defaultSeconds === null) return null;
  return new Date(Date.now() + defaultSeconds * 1000);
}

function getDefaultVirtualKeyScope(
  visibilityOptions: VisibilityOption<VirtualKeyScope>[],
): VirtualKeyScope {
  return (
    visibilityOptions.find((option) => !option.disabled)?.value ?? "personal"
  );
}

function getVirtualKeyVisibilityOptions(params: {
  isAdmin: boolean;
  canReadTeams: boolean;
}): VisibilityOption<VirtualKeyScope>[] {
  const { isAdmin, canReadTeams } = params;

  return [
    {
      value: "personal",
      label: "Personal",
      description: "Only you can view and manage this virtual key",
      icon: User,
    },
    {
      value: "team",
      label: "Team",
      description: "Visible to selected teams",
      icon: Users,
      disabled: !canReadTeams,
      disabledReason: !canReadTeams
        ? "Team sharing is unavailable without team:read permission"
        : undefined,
    },
    {
      value: "org",
      label: "Organization",
      description: "Visible to everyone in the organization",
      icon: Globe,
      disabled: !isAdmin,
      disabledReason: !isAdmin
        ? "You need llmVirtualKey:admin permission to share org-wide"
        : undefined,
    },
  ];
}
