prowler/ui/components/integrations/s3/s3-integration-form.tsx

"use client";

import { zodResolver } from "@hookform/resolvers/zod";
import { ArrowLeftIcon, ArrowRightIcon } from "lucide-react";
import { useSession } from "next-auth/react";
import { useState } from "react";
import type { Control } from "react-hook-form";
import { useForm } from "react-hook-form";

import { createIntegration, updateIntegration } from "@/actions/integrations";
import {
  PROVIDER_TYPE_DATA,
  ProviderTypeIcon,
} from "@/components/icons/providers-badge/provider-type-icon";
import { AWSRoleCredentialsForm } from "@/components/providers/workflow/forms/select-credentials-type/aws/credentials-type/aws-role-credentials-form";
import { Separator } from "@/components/shadcn";
import { useToast } from "@/components/shadcn";
import { CustomInput } from "@/components/shadcn/custom";
import { CustomLink } from "@/components/shadcn/custom/custom-link";
import {
  Form,
  FormControl,
  FormField,
  FormMessage,
} from "@/components/shadcn/form";
import { FormButtons } from "@/components/shadcn/form/form-buttons";
import { EnhancedMultiSelect } from "@/components/shadcn/select/enhanced-multi-select";
import { getAWSCredentialsTemplateLinks } from "@/lib";
import type { AWSCredentialsRole } from "@/types";
import type { IntegrationProps } from "@/types/integrations";
import {
  editS3IntegrationFormSchema,
  s3IntegrationFormSchema,
} from "@/types/integrations";
import type { ProviderProps } from "@/types/providers";

interface S3IntegrationFormProps {
  integration?: IntegrationProps | null;
  providers: ProviderProps[];
  onSuccess: (integrationId?: string, shouldTestConnection?: boolean) => void;
  onCancel: () => void;
  editMode?: "configuration" | "credentials" | null; // null means creating new
}

const getSelectedAWSAccountId = (
  selectedProviderIds: string[],
  providers: ProviderProps[],
): string => {
  for (const providerId of selectedProviderIds) {
    const provider = providers.find(({ id }) => id === providerId);
    const uid = provider?.attributes.uid;
    if (
      provider?.attributes.provider === "aws" &&
      uid &&
      /^\d{12}$/.test(uid)
    ) {
      return uid;
    }
  }

  return "";
};

export const S3IntegrationForm = ({
  integration,
  providers,
  onSuccess,
  onCancel,
  editMode = null,
}: S3IntegrationFormProps) => {
  const { data: session } = useSession();
  const { toast } = useToast();
  const [currentStep, setCurrentStep] = useState(
    editMode === "credentials" ? 1 : 0,
  );
  const isEditing = !!integration;
  const isCreating = !isEditing;
  const isEditingConfig = editMode === "configuration";
  const isEditingCredentials = editMode === "credentials";

  const defaultCredentialsType =
    process.env.NEXT_PUBLIC_IS_CLOUD_ENV === "true"
      ? "aws-sdk-default"
      : "access-secret-key";

  const form = useForm({
    resolver: zodResolver(
      // For credentials editing, use creation schema (all fields required)
      // For config editing, use edit schema (partial updates allowed)
      // For creation, use creation schema
      isEditingCredentials || isCreating
        ? s3IntegrationFormSchema
        : editS3IntegrationFormSchema,
    ),
    defaultValues: {
      integration_type: "amazon_s3" as const,
      bucket_name: integration?.attributes.configuration.bucket_name || "",
      bucket_account_id: "",
      output_directory:
        integration?.attributes.configuration.output_directory || "output",
      providers:
        integration?.relationships?.providers?.data?.map((p) => p.id) || [],
      enabled: integration?.attributes.enabled ?? true,
      credentials_type: defaultCredentialsType,
      aws_access_key_id: "",
      aws_secret_access_key: "",
      aws_session_token: "",
      role_arn: "",
      // External ID always defaults to tenantId, even when editing credentials
      external_id: session?.tenantId || "",
      role_session_name: "",
      session_duration: "",
      show_role_section: false,
    },
  });

  const isLoading = form.formState.isSubmitting;
  const selectedProviderIds = form.watch("providers") || [];
  const bucketAccountIdOverride = form.watch("bucket_account_id")?.trim() || "";
  const derivedBucketAccountId = getSelectedAWSAccountId(
    selectedProviderIds,
    providers,
  );
  const resolvedBucketAccountId =
    bucketAccountIdOverride || derivedBucketAccountId;

  const handleNext = async (e: React.FormEvent) => {
    e.preventDefault();

    // If we're in single-step edit mode, don't advance
    if (isEditingConfig || isEditingCredentials) {
      return;
    }

    // Validate current step fields for creation flow. bucket_account_id is
    // validated here, while its input is visible, so a malformed value surfaces
    // its error instead of silently blocking the step 1 submit.
    const stepFields =
      currentStep === 0
        ? ([
            "bucket_name",
            "output_directory",
            "providers",
            "bucket_account_id",
          ] as const)
        : // Step 1: No required fields since role_arn and external_id are optional
          [];

    const isValid = stepFields.length === 0 || (await form.trigger(stepFields));

    if (!isValid) {
      return;
    }

    if (!resolvedBucketAccountId) {
      form.setError("bucket_account_id", {
        message:
          "Bucket owner account ID is required when no AWS account is selected",
      });
      return;
    }

    form.clearErrors("bucket_account_id");
    setCurrentStep(1);
  };

  const handleBack = () => {
    setCurrentStep(0);
  };

  // Helper function to build credentials object
  const buildCredentials = (values: any) => {
    const credentials: any = {};

    // Don't include credentials_type in the API payload - it's a UI-only field
    // The backend determines credential type based on which fields are present

    // Only include role-related fields if role_arn is provided
    if (values.role_arn && values.role_arn.trim() !== "") {
      credentials.role_arn = values.role_arn;
      credentials.external_id = values.external_id;

      // Optional role fields
      if (values.role_session_name)
        credentials.role_session_name = values.role_session_name;
      if (values.session_duration)
        credentials.session_duration =
          parseInt(values.session_duration, 10) || 3600;
    }

    // Add static credentials if using access-secret-key type
    if (values.credentials_type === "access-secret-key") {
      credentials.aws_access_key_id = values.aws_access_key_id;
      credentials.aws_secret_access_key = values.aws_secret_access_key;
      if (values.aws_session_token)
        credentials.aws_session_token = values.aws_session_token;
    }

    return credentials;
  };

  const buildConfiguration = (values: any, isPartial = false) => {
    const configuration: any = {};

    // For creation mode, include all fields
    if (!isPartial) {
      configuration.bucket_name = values.bucket_name;
      configuration.output_directory = values.output_directory || "output";
    } else {
      // For edit mode, bucket_name and output_directory are treated as a pair
      const originalBucketName =
        integration?.attributes.configuration.bucket_name || "";
      const originalOutputDirectory =
        integration?.attributes.configuration.output_directory || "";

      const bucketNameChanged = values.bucket_name !== originalBucketName;
      const outputDirectoryChanged =
        values.output_directory !== originalOutputDirectory;

      // If either field changed, send both (as a pair)
      if (bucketNameChanged || outputDirectoryChanged) {
        configuration.bucket_name = values.bucket_name;
        configuration.output_directory = values.output_directory || "output";
      }
    }

    return configuration;
  };

  // Helper function to build FormData based on edit mode
  const buildFormData = (values: any) => {
    const formData = new FormData();
    formData.append("integration_type", values.integration_type);

    if (isEditingConfig) {
      const configuration = buildConfiguration(values, true);
      if (Object.keys(configuration).length > 0) {
        formData.append("configuration", JSON.stringify(configuration));
      }
      // Always send providers array, even if empty, to update relationships
      formData.append("providers", JSON.stringify(values.providers || []));
    } else if (isEditingCredentials) {
      const credentials = buildCredentials(values);
      formData.append("credentials", JSON.stringify(credentials));
    } else {
      // Creation mode - send everything
      const configuration = buildConfiguration(values);
      const credentials = buildCredentials(values);

      formData.append("configuration", JSON.stringify(configuration));
      formData.append("credentials", JSON.stringify(credentials));
      formData.append("providers", JSON.stringify(values.providers));
      formData.append("enabled", JSON.stringify(values.enabled ?? true));
    }

    return formData;
  };

  const onSubmit = async (values: any) => {
    const formData = buildFormData(values);

    try {
      let result;
      let shouldTestConnection = false;

      if (isEditing && integration) {
        result = await updateIntegration(integration.id, formData);
        // Test connection if we're editing credentials or configuration (S3 needs both)
        shouldTestConnection = isEditingCredentials || isEditingConfig;
      } else {
        result = await createIntegration(formData);
        // Always test connection for new integrations
        shouldTestConnection = true;
      }

      if ("success" in result) {
        toast({
          title: "Success!",
          description: `S3 integration ${isEditing ? "updated" : "created"} successfully.`,
        });

        // Pass the integration ID and whether to test connection to the success callback
        onSuccess(result.integrationId, shouldTestConnection);
      } else if ("error" in result) {
        const errorMessage = result.error;

        toast({
          variant: "destructive",
          title: "S3 Integration Error",
          description: errorMessage,
        });
      }
    } catch (error) {
      const errorMessage =
        error instanceof Error ? error.message : "An unexpected error occurred";

      toast({
        variant: "destructive",
        title: "Connection Error",
        description: `${errorMessage}. Please check your network connection and try again.`,
      });
    }
  };

  const renderBucketAccountIdField = () => (
    <div className="flex flex-col gap-1">
      <CustomInput
        control={form.control}
        name="bucket_account_id"
        type="text"
        label="Bucket owner account ID"
        labelPlacement="inside"
        placeholder={
          derivedBucketAccountId
            ? `Defaults to ${derivedBucketAccountId}`
            : "12-digit AWS account ID"
        }
        variant="bordered"
        isRequired={false}
      />
      <p className="text-text-neutral-tertiary text-xs">
        {derivedBucketAccountId
          ? `Leave empty to use selected AWS account ${derivedBucketAccountId}, or enter another bucket owner account ID.`
          : "Required because the selected provider does not identify the AWS account that owns the bucket."}
      </p>
    </div>
  );

  const renderStepContent = () => {
    // If editing credentials, show only credentials form
    if (isEditingCredentials || currentStep === 1) {
      const bucketName = form.getValues("bucket_name") || "";
      const externalId =
        form.getValues("external_id") || session?.tenantId || "";
      const templateLinks = getAWSCredentialsTemplateLinks(
        externalId,
        bucketName,
        "amazon_s3",
        resolvedBucketAccountId,
      );

      return (
        <div className="flex flex-col gap-4">
          {isEditingCredentials && renderBucketAccountIdField()}
          <AWSRoleCredentialsForm
            control={form.control as unknown as Control<AWSCredentialsRole>}
            setValue={form.setValue as any}
            externalId={externalId}
            templateLinks={templateLinks}
            type="integrations"
            integrationType="amazon_s3"
          />
        </div>
      );
    }

    // Show configuration step (step 0 or editing configuration)
    if (isEditingConfig || currentStep === 0) {
      const providerOptions = providers.map((provider) => {
        const providerType = provider.attributes.provider;
        return {
          value: provider.id,
          label: provider.attributes.alias || provider.attributes.uid,
          icon:
            providerType in PROVIDER_TYPE_DATA ? (
              <ProviderTypeIcon type={providerType} size={20} />
            ) : undefined,
          description: provider.attributes.connection.connected
            ? "Connected"
            : "Disconnected",
        };
      });

      return (
        <>
          {/* Provider Selection */}
          <div className="flex flex-col gap-4">
            <FormField
              control={form.control}
              name="providers"
              render={({ field }) => (
                <>
                  <FormControl>
                    <EnhancedMultiSelect
                      options={providerOptions}
                      onValueChange={field.onChange}
                      defaultValue={field.value || []}
                      placeholder="Select providers to integrate with"
                      searchable={true}
                      maxCount={1}
                    />
                  </FormControl>
                  <FormMessage className="text-text-error max-w-full text-xs" />
                </>
              )}
            />
          </div>

          <Separator />

          {/* S3 Configuration */}
          <div className="flex flex-col gap-4">
            <CustomInput
              control={form.control}
              name="bucket_name"
              type="text"
              label="Bucket name"
              labelPlacement="inside"
              placeholder="my-security-findings-bucket"
              variant="bordered"
              isRequired
            />

            <CustomInput
              control={form.control}
              name="output_directory"
              type="text"
              label="Output directory"
              labelPlacement="inside"
              placeholder="output"
              variant="bordered"
              isRequired
            />

            {!isEditingConfig && renderBucketAccountIdField()}
          </div>
        </>
      );
    }

    return null;
  };

  const renderStepButtons = () => {
    // Single edit mode (configuration or credentials)
    if (isEditingConfig || isEditingCredentials) {
      const updateText = isEditingConfig
        ? "Update Configuration"
        : "Update Credentials";
      const loadingText = isEditingConfig
        ? "Updating Configuration..."
        : "Updating Credentials...";

      return (
        <FormButtons
          setIsOpen={() => {}}
          onCancel={onCancel}
          submitText={updateText}
          cancelText="Cancel"
          loadingText={loadingText}
          isDisabled={isLoading}
        />
      );
    }

    // Creation flow - step 0
    if (currentStep === 0) {
      return (
        <FormButtons
          setIsOpen={() => {}}
          onCancel={onCancel}
          submitText="Next"
          cancelText="Cancel"
          loadingText="Processing..."
          isDisabled={isLoading}
          rightIcon={<ArrowRightIcon size={24} />}
        />
      );
    }

    // Creation flow - step 1 (final step)
    return (
      <FormButtons
        setIsOpen={() => {}}
        onCancel={handleBack}
        submitText="Create Integration"
        cancelText="Back"
        loadingText="Creating..."
        leftIcon={<ArrowLeftIcon size={24} />}
        isDisabled={isLoading}
      />
    );
  };

  return (
    <Form {...form}>
      <form
        onSubmit={
          // For edit modes, always submit
          isEditingConfig || isEditingCredentials
            ? form.handleSubmit(onSubmit)
            : // For creation flow, handle step logic
              currentStep === 0
              ? handleNext
              : form.handleSubmit(onSubmit)
        }
        className="flex flex-col gap-6"
      >
        <div className="flex flex-col gap-4">
          <div className="flex flex-col items-start gap-2 sm:flex-row sm:items-center">
            <p className="text-text-neutral-tertiary flex items-center gap-2 text-sm">
              Need help configuring your Amazon S3 integration?
            </p>
            <CustomLink
              href="https://docs.prowler.com/projects/prowler-open-source/en/latest/tutorials/prowler-app-s3-integration/"
              target="_blank"
              size="sm"
            >
              Read the docs
            </CustomLink>
          </div>
          {renderStepContent()}
        </div>
        {renderStepButtons()}
      </form>
    </Form>
  );
};