prowler/ui/components/providers/workflow/forms/connect-account-form.tsx

"use client";

import { zodResolver } from "@hookform/resolvers/zod";
import { ChevronLeftIcon, ChevronRightIcon, Loader2 } from "lucide-react";
import { useRouter } from "next/navigation";
import { Dispatch, SetStateAction, useEffect, useRef, useState } from "react";
import { useForm, UseFormReturn } from "react-hook-form";

import { addProvider, updateProvider } from "@/actions/providers/providers";
import { addRegistryProvider } from "@/actions/providers/registry-provider";
import { AzureMethodSelector } from "@/components/providers/organizations/azure-method-selector";
import { GcpMethodSelector } from "@/components/providers/organizations/gcp-method-selector";
import { WizardInputField } from "@/components/providers/workflow/forms/fields";
import { ProviderTitleDocs } from "@/components/providers/workflow/provider-title-docs";
import {
  RegistryUpgradePrompt,
  RegistrySetupPrompt,
} from "@/components/registry/registry-upgrade-prompt";
import { Button, useToast } from "@/components/shadcn";
import { Alert, AlertDescription, AlertTitle } from "@/components/shadcn/alert";
import { Form } from "@/components/shadcn/form";
import { ProviderCredentialFields } from "@/lib/provider-credentials/provider-credential-fields";
import { REGISTRY_PROVIDER_DISCOVERY } from "@/lib/registry/provider-options";
import { isRegistryEnabled } from "@/lib/shared/env";
import {
  createAddProviderFormSchema,
  AddProviderFormValues,
  ApiError,
  KnownProviderType,
  ProviderType,
} from "@/types";
import {
  ORGANIZATION_TYPE,
  OrgFlowType,
  toOrgFlowType,
} from "@/types/organizations";
import { isKnownProviderType } from "@/types/providers";

import { RadioGroupProvider } from "../../radio-group-provider";

import { useRegistryProviderDiscovery } from "./use-registry-provider-discovery";

export type FormValues = AddProviderFormValues;

export interface ConnectAccountSuccessData {
  id: string;
  providerType: ProviderType;
  uid: string;
  alias: string | null;
}

/**
 * Provider types that offer an organization-onboarding method choice: exactly the
 * ones with an onboarding flow, so a new flow type cannot miss the fork. AWS is the
 * exception: the wizard's own AWS step hosts its single-account/organization switch.
 */
function providerHasOrgMethod(
  providerType: ProviderType | undefined,
): providerType is OrgFlowType {
  return providerType !== "aws" && toOrgFlowType(providerType) !== undefined;
}

interface ConnectAccountFormProps {
  onSuccess?: (data: ConnectAccountSuccessData) => void;
  onSelectOrganizations?: (orgType: OrgFlowType) => void;
  onProviderTypeChange?: (providerType: ProviderType | null) => void;
  formId?: string;
  hideNavigation?: boolean;
  onUiStateChange?: (state: {
    showBack: boolean;
    showAction: boolean;
    actionLabel: string;
    actionDisabled: boolean;
    isLoading: boolean;
  }) => void;
  onBackHandlerChange?: (handler: () => void) => void;
}

// Helper function for labels and placeholders
const getProviderFieldDetails = (providerType?: ProviderType) => {
  switch (providerType) {
    case "aws":
      return {
        label: "Account ID",
        placeholder: "e.g. 123456789012",
      };
    case "gcp":
      return {
        label: "Project ID",
        placeholder: "e.g. my-gcp-project",
      };
    case "azure":
      return {
        label: "Subscription ID",
        placeholder: "e.g. fc94207a-d396-4a14-a7fd-12ab34cd56ef",
      };
    case "kubernetes":
      return {
        label: "Kubernetes Context",
        placeholder: "e.g. my-cluster-context",
      };
    case "m365":
      return {
        label: "Domain ID",
        placeholder: "e.g. your-domain.onmicrosoft.com",
      };
    case "github":
      return {
        label: "Username/Organization",
        placeholder: "e.g. username or organization-name",
      };
    case "iac":
      return {
        label: "Repository URL",
        placeholder: "e.g. https://github.com/user/repo",
      };
    case "image":
      return {
        label: "Registry URL",
        placeholder: "e.g. 123456789012.dkr.ecr.us-east-1.amazonaws.com",
      };
    case "oraclecloud":
      return {
        label: "Tenancy OCID",
        placeholder: "e.g. ocid1.tenancy.oc1..aaaaaaa...",
      };
    case "mongodbatlas":
      return {
        label: "Organization ID",
        placeholder: "e.g. 5f43a8c4e1234567890abcde",
      };
    case "alibabacloud":
      return {
        label: "Account ID",
        placeholder: "e.g. 1234567890123456",
      };
    case "cloudflare":
      return {
        label: "Account ID",
        placeholder: "e.g. a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4",
      };
    case "openstack":
      return {
        label: "Project ID",
        placeholder: "e.g. a1b2c3d4-e5f6-7890-abcd-ef1234567890",
      };
    case "googleworkspace":
      return {
        label: "Customer ID",
        placeholder: "e.g. C01234abc",
      };
    case "vercel":
      return {
        label: "Team ID",
        placeholder: "e.g. team_xxxxxxxxxxxxxxxxxxxxxxxx",
      };
    case "okta":
      return {
        label: "Org Domain",
        placeholder: "e.g. your-org.okta.com",
      };
    default:
      return {
        label: "Provider UID",
        placeholder: "Enter the provider UID",
      };
  }
};

function applyBackStep({
  prevStep,
  method,
  providerType,
  form,
  setPrevStep,
  setMethod,
}: {
  prevStep: number;
  method: "single" | null;
  providerType: ProviderType | undefined;
  form: Pick<UseFormReturn<FormValues>, "setValue" | "clearErrors">;
  setPrevStep: Dispatch<SetStateAction<number>>;
  setMethod: Dispatch<SetStateAction<"single" | null>>;
}) {
  // With a method choice, "Back" from the single account/project form returns to
  // the method selector rather than the provider picker.
  if (
    prevStep === 2 &&
    method === "single" &&
    providerHasOrgMethod(providerType)
  ) {
    setMethod(null);
    form.setValue("providerUid", "", { shouldValidate: false });
    form.setValue("providerAlias", "", { shouldValidate: false });
    return;
  }

  setPrevStep((prev) => prev - 1);
  if (prevStep === 2) {
    form.setValue("providerType", undefined as unknown as KnownProviderType, {
      shouldValidate: false,
    });
    setMethod(null);
  }
  form.setValue("providerUid", "", { shouldValidate: false });
  form.setValue("providerAlias", "", { shouldValidate: false });
  form.clearErrors();
}

export const ConnectAccountForm = ({
  onSuccess,
  onSelectOrganizations,
  onProviderTypeChange,
  formId,
  hideNavigation = false,
  onUiStateChange,
  onBackHandlerChange,
}: ConnectAccountFormProps) => {
  const { toast } = useToast();
  const [prevStep, setPrevStep] = useState(1);
  const [method, setMethod] = useState<"single" | null>(null);
  const router = useRouter();

  const {
    result: registryDiscovery,
    isRetryingDiscovery,
    retryDiscovery,
  } = useRegistryProviderDiscovery();
  const registryOptions =
    registryDiscovery?.status === REGISTRY_PROVIDER_DISCOVERY.READY
      ? registryDiscovery.options
      : [];
  // The deployment flag is known up front, so the tabs render before discovery.
  const registryAvailable =
    isRegistryEnabled() &&
    registryDiscovery?.status !== REGISTRY_PROVIDER_DISCOVERY.ACCESS_DENIED;
  const registryError =
    registryDiscovery?.status === REGISTRY_PROVIDER_DISCOVERY.ERROR ||
    registryDiscovery?.status === REGISTRY_PROVIDER_DISCOVERY.UNKNOWN;
  const [providerError, setProviderError] = useState<string | null>(null);
  const submitting = useRef(false);
  const createdAccount = useRef<ConnectAccountSuccessData | null>(null);

  const formSchema = createAddProviderFormSchema(
    registryOptions.map((option) => option.type),
  );

  const form = useForm<FormValues>({
    resolver: zodResolver(formSchema),
    defaultValues: {
      providerType: undefined,
      providerUid: "",
      providerAlias: "",
    },
    mode: "onChange",
    reValidateMode: "onChange",
  });

  const providerType = form.watch("providerType");
  const providerUid = form.watch("providerUid");
  const providerFieldDetails = getProviderFieldDetails(providerType);

  const isLoading = form.formState.isSubmitting;

  const onSubmitClient = async (values: FormValues) => {
    if (submitting.current) return;
    const existingAccount =
      createdAccount.current?.providerType === values.providerType &&
      createdAccount.current.uid === values.providerUid
        ? createdAccount.current
        : null;
    if (
      existingAccount &&
      (existingAccount.alias ?? "") === (values.providerAlias?.trim() ?? "") &&
      onSuccess
    ) {
      onSuccess(existingAccount);
      return;
    }
    submitting.current = true;
    setProviderError(null);
    const formValues = { ...values };

    const formData = new FormData();
    Object.entries(formValues).forEach(
      ([key, value]) => value !== undefined && formData.append(key, value),
    );

    try {
      let data;
      if (existingAccount) {
        const update = new FormData();
        update.set(ProviderCredentialFields.PROVIDER_ID, existingAccount.id);
        update.set(
          ProviderCredentialFields.PROVIDER_ALIAS,
          values.providerAlias?.trim() ?? "",
        );
        data = await updateProvider(update);
      } else {
        data = await (isKnownProviderType(values.providerType)
          ? addProvider(formData)
          : addRegistryProvider(formData));
      }

      if (data?.errors && data.errors.length > 0) {
        data.errors.forEach((error: ApiError) => {
          const errorMessage = error.detail;
          const pointer = error.source?.pointer;

          switch (pointer) {
            case "/data/attributes/provider":
              // Provider selection is hidden here; keep failures visible and
              // retryable when availability changes without editing the form.
              setProviderError(errorMessage);
              break;
            case "/data/attributes/uid":
            case "/data/attributes/__all__":
              form.setError("providerUid", {
                type: "server",
                message: errorMessage,
              });
              break;
            case "/data/attributes/alias":
              form.setError("providerAlias", {
                type: "server",
                message: errorMessage,
              });
              break;
            default:
              toast({
                variant: "destructive",
                title: "Oops! Something went wrong",
                description: errorMessage,
              });
          }
        });
        return;
      } else {
        const {
          id,
          attributes: { provider: createdProviderType, uid, alias },
        } = data.data;

        if (onSuccess) {
          createdAccount.current = {
            id,
            providerType: createdProviderType,
            uid: uid || values.providerUid,
            alias: alias ?? values.providerAlias ?? null,
          };
          onSuccess(createdAccount.current);
          return;
        }

        router.push("/providers");
      }
    } catch (error: unknown) {
      console.error("Error during submission:", error);
      toast({
        variant: "destructive",
        title: "Submission Error",
        description:
          error instanceof Error
            ? error.message
            : "Something went wrong. Please try again.",
      });
    } finally {
      submitting.current = false;
    }
  };

  const handleBackStep = () => {
    setProviderError(null);
    applyBackStep({
      prevStep,
      method,
      providerType,
      form,
      setPrevStep,
      setMethod,
    });
  };

  useEffect(() => {
    if (providerType) {
      setPrevStep(2);
    }
  }, [providerType]);

  useEffect(() => {
    onProviderTypeChange?.(providerType ?? null);
  }, [onProviderTypeChange, providerType]);

  useEffect(() => {
    onBackHandlerChange?.(() => {
      setProviderError(null);
      applyBackStep({
        prevStep,
        method,
        providerType,
        form,
        setPrevStep,
        setMethod,
      });
    });
  }, [onBackHandlerChange, prevStep, method, providerType, form]);

  // Providers with a method choice reach the UID form only through "single".
  const showUidForm =
    !providerHasOrgMethod(providerType) || method === "single";

  useEffect(() => {
    const canSubmit =
      prevStep === 2 &&
      showUidForm &&
      providerUid.trim().length > 0 &&
      form.formState.isValid;

    onUiStateChange?.({
      showBack: prevStep === 2,
      showAction: prevStep === 2 && showUidForm,
      actionLabel: isLoading ? "Creating provider..." : "Next",
      actionDisabled: !canSubmit || isLoading,
      isLoading,
    });
  }, [
    showUidForm,
    form.formState.isValid,
    isLoading,
    onUiStateChange,
    prevStep,
    providerUid,
  ]);

  return (
    <Form {...form}>
      <form
        id={formId}
        onSubmit={form.handleSubmit(onSubmitClient)}
        className="flex flex-col gap-4"
      >
        {/* Step 1: Provider selection */}
        {prevStep === 1 && (
          <div>
            <RadioGroupProvider
              registryAvailable={registryAvailable}
              registryLoading={registryAvailable && registryDiscovery === null}
              registryNotice={
                registryDiscovery?.status ===
                REGISTRY_PROVIDER_DISCOVERY.SUBSCRIPTION_REQUIRED ? (
                  <RegistryUpgradePrompt />
                ) : registryDiscovery?.status ===
                  REGISTRY_PROVIDER_DISCOVERY.ONBOARDING ? (
                  <RegistrySetupPrompt />
                ) : (
                  registryError && (
                    <Alert variant="warning">
                      <AlertTitle>
                        Registry providers could not be loaded
                      </AlertTitle>
                      <AlertDescription>
                        Built-in providers are available. Check the Registry
                        connection and try again.
                        {/* aria-disabled, not disabled: the pressed button keeps focus. */}
                        <Button
                          type="button"
                          variant="outline"
                          aria-disabled={isRetryingDiscovery}
                          onClick={retryDiscovery}
                        >
                          {isRetryingDiscovery
                            ? "Retrying…"
                            : "Retry Registry providers"}
                        </Button>
                      </AlertDescription>
                    </Alert>
                  )
                )
              }
              registryOptions={registryOptions}
              control={form.control}
              isInvalid={!!form.formState.errors.providerType}
              errorMessage={form.formState.errors.providerType?.message}
            />
          </div>
        )}
        {/* Step 2: Azure method selector (before choosing a method) */}
        {prevStep === 2 && providerType === "azure" && method === null && (
          <>
            <ProviderTitleDocs providerType={providerType} />
            <AzureMethodSelector
              onSelectSingle={() => setMethod("single")}
              onSelectOrganizations={() =>
                onSelectOrganizations?.(ORGANIZATION_TYPE.AZURE)
              }
            />
          </>
        )}
        {/* Step 2: GCP method selector (before choosing a method) */}
        {prevStep === 2 && providerType === "gcp" && method === null && (
          <>
            <ProviderTitleDocs providerType={providerType} />
            <GcpMethodSelector
              onSelectSingle={() => setMethod("single")}
              onSelectOrganizations={() =>
                onSelectOrganizations?.(ORGANIZATION_TYPE.GCP)
              }
            />
          </>
        )}
        {/* Step 2: UID, alias form (providers without a method choice, or the
            single account/subscription/project method) */}
        {prevStep === 2 && showUidForm && (
          <>
            <ProviderTitleDocs providerType={providerType} />
            {providerError && (
              <Alert variant="destructive">
                <AlertTitle>Unable to create provider</AlertTitle>
                <AlertDescription>{providerError}</AlertDescription>
              </Alert>
            )}
            <WizardInputField
              control={form.control}
              name="providerUid"
              type="text"
              label={providerFieldDetails.label}
              labelPlacement="inside"
              placeholder={providerFieldDetails.placeholder}
              variant="bordered"
              isRequired
            />
            <WizardInputField
              control={form.control}
              name="providerAlias"
              type="text"
              label="Provider alias (optional)"
              labelPlacement="inside"
              placeholder="Enter the provider alias"
              variant="bordered"
              isRequired={false}
            />
          </>
        )}
        {!hideNavigation && (
          <div className="flex w-full justify-end gap-4">
            {prevStep === 2 && (
              <Button
                type="button"
                variant="ghost"
                size="lg"
                onClick={handleBackStep}
                disabled={isLoading}
              >
                {!isLoading && <ChevronLeftIcon size={24} />}
                Back
              </Button>
            )}
            {prevStep === 2 && showUidForm && (
              <Button
                type="submit"
                variant="default"
                size="lg"
                disabled={isLoading}
                aria-busy={isLoading || undefined}
              >
                {isLoading ? (
                  <Loader2 aria-hidden className="animate-spin" />
                ) : (
                  <ChevronRightIcon size={24} />
                )}
                {isLoading ? "Creating provider..." : "Next"}
              </Button>
            )}
          </div>
        )}
      </form>
    </Form>
  );
};