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, useState } from "react";
import { useForm, UseFormReturn } from "react-hook-form";
import { z } from "zod";

import { addProvider } from "@/actions/providers/providers";
import { AwsMethodSelector } from "@/components/providers/organizations/aws-method-selector";
import { WizardInputField } from "@/components/providers/workflow/forms/fields";
import { ProviderTitleDocs } from "@/components/providers/workflow/provider-title-docs";
import { Button } from "@/components/shadcn";
import { useToast } from "@/components/ui";
import { Form } from "@/components/ui/form";
import { addProviderFormSchema, ApiError, ProviderType } from "@/types";

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

export type FormValues = z.infer<typeof addProviderFormSchema>;

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

interface ConnectAccountFormProps {
  onSuccess?: (data: ConnectAccountSuccessData) => void;
  onSelectOrganizations?: () => 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 "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",
      };
    default:
      return {
        label: "Provider UID",
        placeholder: "Enter the provider UID",
      };
  }
};

function applyBackStep({
  prevStep,
  awsMethod,
  form,
  setPrevStep,
  setAwsMethod,
}: {
  prevStep: number;
  awsMethod: "single" | null;
  form: Pick<UseFormReturn<FormValues>, "setValue">;
  setPrevStep: Dispatch<SetStateAction<number>>;
  setAwsMethod: Dispatch<SetStateAction<"single" | null>>;
}) {
  // If in UID form after choosing single, go back to method selector
  if (prevStep === 2 && awsMethod === "single") {
    setAwsMethod(null);
    form.setValue("providerUid", "");
    form.setValue("providerAlias", "");
    return;
  }

  setPrevStep((prev) => prev - 1);
  // Deselect the providerType if the user is going back to the first step
  if (prevStep === 2) {
    form.setValue("providerType", undefined as unknown as ProviderType);
    setAwsMethod(null);
  }
  // Reset the providerUid and providerAlias fields when going back
  form.setValue("providerUid", "");
  form.setValue("providerAlias", "");
}

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

  const formSchema = addProviderFormSchema;

  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) => {
    const formValues = { ...values };

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

    try {
      const data = await addProvider(formData);

      if (data?.errors && data.errors.length > 0) {
        // Handle server-side validation errors
        data.errors.forEach((error: ApiError) => {
          const errorMessage = error.detail;
          const pointer = error.source?.pointer;

          switch (pointer) {
            case "/data/attributes/provider":
              form.setError("providerType", {
                type: "server",
                message: 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 {
        // Go to the next step after successful submission
        const {
          id,
          attributes: { provider: createdProviderType, uid, alias },
        } = data.data;

        if (onSuccess) {
          onSuccess({
            id,
            providerType: createdProviderType,
            uid: uid || values.providerUid,
            alias: alias ?? values.providerAlias ?? null,
          });
          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.",
      });
    }
  };

  const handleBackStep = () => {
    applyBackStep({
      prevStep,
      awsMethod,
      form,
      setPrevStep,
      setAwsMethod,
    });
  };

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

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

  useEffect(() => {
    onBackHandlerChange?.(() => {
      applyBackStep({
        prevStep,
        awsMethod,
        form,
        setPrevStep,
        setAwsMethod,
      });
    });
  }, [onBackHandlerChange, prevStep, awsMethod, form]);

  useEffect(() => {
    const canSubmit =
      prevStep === 2 &&
      (providerType !== "aws" || awsMethod === "single") &&
      providerUid.trim().length > 0 &&
      form.formState.isValid;

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

  return (
    <Form {...form}>
      <form
        id={formId}
        onSubmit={form.handleSubmit(onSubmitClient)}
        className="flex flex-col gap-4"
      >
        {/* Step 1: Provider selection */}
        {prevStep === 1 && (
          <RadioGroupProvider
            control={form.control}
            isInvalid={!!form.formState.errors.providerType}
            errorMessage={form.formState.errors.providerType?.message}
          />
        )}
        {/* Step 2: AWS method selector (only for AWS, before choosing method) */}
        {prevStep === 2 && providerType === "aws" && awsMethod === null && (
          <>
            <ProviderTitleDocs providerType={providerType} />
            <AwsMethodSelector
              onSelectSingle={() => setAwsMethod("single")}
              onSelectOrganizations={() => {
                onSelectOrganizations?.();
              }}
            />
          </>
        )}
        {/* Step 2: UID, alias form (non-AWS or AWS single account) */}
        {prevStep === 2 &&
          (providerType !== "aws" || awsMethod === "single") && (
            <>
              <ProviderTitleDocs providerType={providerType} />
              <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 &&
              (providerType !== "aws" || awsMethod === "single") && (
                <Button
                  type="submit"
                  variant="default"
                  size="lg"
                  disabled={isLoading}
                >
                  {isLoading ? (
                    <Loader2 className="animate-spin" />
                  ) : (
                    <ChevronRightIcon size={24} />
                  )}
                  {isLoading ? "Loading" : "Next"}
                </Button>
              )}
          </div>
        )}
      </form>
    </Form>
  );
};