prowler/ui/components/providers/wizard/hooks/use-provider-wizard-controller.ts

"use client";

import { useRouter } from "next/navigation";
import { useEffect, useRef, useState } from "react";

import { DOCS_URLS, getProviderHelpText } from "@/lib/external-urls";
import {
  dispatchProviderFunnel,
  PROVIDER_FUNNEL_METHOD,
  PROVIDER_FUNNEL_STEP,
} from "@/lib/provider-funnel/provider-funnel-events";
import { isCloud } from "@/lib/shared/env";
import { endActiveTour } from "@/lib/tours/use-driver-tour";
import { useOnboardingCheckpointStore } from "@/store/onboarding-checkpoint";
import { useOrgSetupStore } from "@/store/organizations/store";
import { useProviderWizardStore } from "@/store/provider-wizard/store";
import {
  ORG_SETUP_PHASE,
  ORG_WIZARD_STEP,
  ORGANIZATION_TYPE,
  OrgFlowType,
  OrgSetupPhase,
  OrgWizardStep,
} from "@/types/organizations";
import {
  PROVIDER_WIZARD_MODE,
  PROVIDER_WIZARD_STEP,
  ProviderWizardStep,
} from "@/types/provider-wizard";
import { ProviderType } from "@/types/providers";

import { getProviderWizardModalTitle } from "../provider-wizard-modal.utils";
import {
  WIZARD_FOOTER_ACTION_TYPE,
  WizardFooterConfig,
} from "../steps/footer-controls";
import type { OrgWizardInitialData, ProviderWizardInitialData } from "../types";

const WIZARD_VARIANT = {
  PROVIDER: "provider",
  ORGANIZATIONS: "organizations",
} as const;

type WizardVariant = (typeof WIZARD_VARIANT)[keyof typeof WIZARD_VARIANT];

const ORG_DOCS_URL = {
  [ORGANIZATION_TYPE.AWS]: DOCS_URLS.AWS_ORGANIZATIONS,
  [ORGANIZATION_TYPE.AZURE]: DOCS_URLS.AZURE_ORGANIZATIONS,
  [ORGANIZATION_TYPE.GCP]: DOCS_URLS.GCP_ORGANIZATIONS,
} as const satisfies Record<OrgFlowType, string>;

// Stable names for the abandonment signal; the numeric step ids are not a contract.
const PROVIDER_STEP_NAME = {
  [PROVIDER_WIZARD_STEP.CONNECT]: "connect",
  [PROVIDER_WIZARD_STEP.CREDENTIALS]: "credentials",
  [PROVIDER_WIZARD_STEP.TEST]: "test",
  [PROVIDER_WIZARD_STEP.LAUNCH]: "launch",
} as const satisfies Record<ProviderWizardStep, string>;

const ORG_STEP_NAME = {
  [ORG_WIZARD_STEP.SETUP]: "organizations_setup",
  [ORG_WIZARD_STEP.VALIDATE]: "organizations_validate",
  [ORG_WIZARD_STEP.LAUNCH]: "organizations_launch",
} as const satisfies Record<OrgWizardStep, string>;

const EMPTY_FOOTER_CONFIG: WizardFooterConfig = {
  showBack: false,
  backLabel: "Back",
  showSecondaryAction: false,
  secondaryActionLabel: "",
  secondaryActionVariant: "outline",
  secondaryActionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
  showAction: false,
  actionLabel: "Next",
  actionType: WIZARD_FOOTER_ACTION_TYPE.BUTTON,
};

interface UseProviderWizardControllerProps {
  open: boolean;
  onOpenChange: (open: boolean) => void;
  initialData?: ProviderWizardInitialData;
  orgInitialData?: OrgWizardInitialData;
  // When false, skips post-close router.refresh() — caller relies on revalidatePath instead.
  refreshOnClose?: boolean;
}

export function useProviderWizardController({
  open,
  onOpenChange,
  initialData,
  orgInitialData,
  refreshOnClose = true,
}: UseProviderWizardControllerProps) {
  const router = useRouter();
  const initialProviderId = initialData?.providerId ?? null;
  const initialProviderType = initialData?.providerType ?? null;
  const initialProviderUid = initialData?.providerUid ?? null;
  const initialProviderAlias = initialData?.providerAlias ?? null;
  const initialSecretId = initialData?.secretId ?? null;
  const initialVia = initialData?.via ?? null;
  const initialMode = initialData?.mode ?? null;
  const hasHydratedForCurrentOpenRef = useRef(false);
  const [wizardVariant, setWizardVariant] = useState<WizardVariant>(
    WIZARD_VARIANT.PROVIDER,
  );
  const [currentStep, setCurrentStep] = useState<ProviderWizardStep>(
    PROVIDER_WIZARD_STEP.CONNECT,
  );
  const [orgCurrentStep, setOrgCurrentStep] = useState<OrgWizardStep>(
    ORG_WIZARD_STEP.SETUP,
  );
  const [footerConfig, setFooterConfig] =
    useState<WizardFooterConfig>(EMPTY_FOOTER_CONFIG);
  const [providerTypeHint, setProviderTypeHint] = useState<ProviderType | null>(
    null,
  );
  const [orgSetupPhase, setOrgSetupPhase] = useState<OrgSetupPhase>(
    ORG_SETUP_PHASE.DETAILS,
  );

  const {
    reset: resetProviderWizard,
    setProvider,
    setVia,
    setSecretId,
    setMode,
    mode,
    providerType,
    via,
  } = useProviderWizardStore();
  const {
    reset: resetOrgWizard,
    setOrganization,
    setOrganizationType,
    organizationType,
  } = useOrgSetupStore();

  useEffect(() => {
    if (!open) {
      hasHydratedForCurrentOpenRef.current = false;
      return;
    }

    if (hasHydratedForCurrentOpenRef.current) {
      return;
    }
    hasHydratedForCurrentOpenRef.current = true;

    if (orgInitialData) {
      setWizardVariant(WIZARD_VARIANT.ORGANIZATIONS);
      resetOrgWizard();
      setOrganizationType(orgInitialData.organizationType);
      setOrganization(
        orgInitialData.organizationId,
        orgInitialData.organizationName,
        orgInitialData.externalId,
      );
      setOrgCurrentStep(orgInitialData.targetStep);
      setOrgSetupPhase(orgInitialData.targetPhase);
      setFooterConfig(EMPTY_FOOTER_CONFIG);
      setProviderTypeHint(null);
      return;
    }

    if (initialProviderId && initialProviderType && initialProviderUid) {
      setWizardVariant(WIZARD_VARIANT.PROVIDER);
      setProvider({
        id: initialProviderId,
        type: initialProviderType,
        uid: initialProviderUid,
        alias: initialProviderAlias,
      });
      setVia(initialVia);
      setSecretId(initialSecretId);
      setMode(
        initialMode ||
          (initialSecretId
            ? PROVIDER_WIZARD_MODE.UPDATE
            : PROVIDER_WIZARD_MODE.ADD),
      );
      setCurrentStep(PROVIDER_WIZARD_STEP.CREDENTIALS);
      setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
      setFooterConfig(EMPTY_FOOTER_CONFIG);
      setProviderTypeHint(initialProviderType);
      setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
      return;
    }

    resetProviderWizard();
    resetOrgWizard();
    setWizardVariant(WIZARD_VARIANT.PROVIDER);
    setCurrentStep(PROVIDER_WIZARD_STEP.CONNECT);
    setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
    setFooterConfig(EMPTY_FOOTER_CONFIG);
    setProviderTypeHint(null);
    setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
  }, [
    initialMode,
    initialProviderAlias,
    initialProviderId,
    initialProviderType,
    initialProviderUid,
    initialSecretId,
    initialVia,
    open,
    orgInitialData,
    resetOrgWizard,
    resetProviderWizard,
    setMode,
    setOrganization,
    setOrganizationType,
    setProvider,
    setSecretId,
    setVia,
  ]);

  const isOrgDirectEntry = Boolean(orgInitialData);
  // Opened on an existing account's credentials, so the one-step AWS flow is not
  // in play. Same three fields the hydration above requires to start on CREDENTIALS.
  const isDirectCredentialsEntry = Boolean(
    initialProviderId && initialProviderType && initialProviderUid,
  );

  const handleClose = () => {
    // Closing the wizard at any point ends the add-provider tour; the checkpoint
    // logic below still drives the handoff to scans. No-op off-onboarding.
    endActiveTour();

    // Read providerId before reset clears it — non-null means a provider was connected.
    const connectedProviderId = useProviderWizardStore.getState().providerId;

    dispatchProviderFunnel({
      step: PROVIDER_FUNNEL_STEP.WIZARD_CLOSED,
      lastStep:
        wizardVariant === WIZARD_VARIANT.PROVIDER
          ? PROVIDER_STEP_NAME[currentStep]
          : ORG_STEP_NAME[orgCurrentStep],
      providerCreated: connectedProviderId !== null,
    });

    resetProviderWizard();
    resetOrgWizard();
    setWizardVariant(WIZARD_VARIANT.PROVIDER);
    setCurrentStep(PROVIDER_WIZARD_STEP.CONNECT);
    setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
    setFooterConfig(EMPTY_FOOTER_CONFIG);
    setProviderTypeHint(null);
    setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
    onOpenChange(false);

    // Cloud-only; only fires if the store is armed (user started onboarding).
    if (isCloud()) {
      useOnboardingCheckpointStore.getState().requestOpenOnWizardClose({
        providerConnected: connectedProviderId !== null,
      });
    }

    if (refreshOnClose) {
      router.refresh();
    }
  };

  const handleDialogOpenChange = (nextOpen: boolean) => {
    if (nextOpen) {
      onOpenChange(true);
      return;
    }
    handleClose();
  };

  const handleTestSuccess = () => {
    if (
      useProviderWizardStore.getState().mode === PROVIDER_WIZARD_MODE.UPDATE
    ) {
      handleClose();
      return;
    }
    setCurrentStep(PROVIDER_WIZARD_STEP.LAUNCH);
  };

  const openOrganizationsFlow = (
    orgType: OrgFlowType = ORGANIZATION_TYPE.AWS,
  ) => {
    // Organizations diverges from the credentials path the tour guides toward; end
    // it so it doesn't dangle on a step that no longer fits. No-op off-onboarding.
    endActiveTour();
    dispatchProviderFunnel({
      step: PROVIDER_FUNNEL_STEP.METHOD_SELECTED,
      providerType: orgType,
      method: PROVIDER_FUNNEL_METHOD.ORGANIZATION,
    });
    resetOrgWizard();
    setOrganizationType(orgType);
    setWizardVariant(WIZARD_VARIANT.ORGANIZATIONS);
    setOrgCurrentStep(ORG_WIZARD_STEP.SETUP);
    setFooterConfig(EMPTY_FOOTER_CONFIG);
    setProviderTypeHint(null);
    setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
  };

  const backToProviderFlow = () => {
    // The AWS organization flow is entered from the AWS connect step's tabs, so
    // going back lands on that step again instead of the provider picker.
    const cameFromAwsConnect = organizationType === ORGANIZATION_TYPE.AWS;
    resetOrgWizard();
    setWizardVariant(WIZARD_VARIANT.PROVIDER);
    setCurrentStep(PROVIDER_WIZARD_STEP.CONNECT);
    setFooterConfig(EMPTY_FOOTER_CONFIG);
    setProviderTypeHint(cameFromAwsConnect ? "aws" : null);
    setOrgSetupPhase(ORG_SETUP_PHASE.DETAILS);
  };

  const isProviderFlow = wizardVariant === WIZARD_VARIANT.PROVIDER;
  const docsLink = isProviderFlow
    ? getProviderHelpText(
        providerTypeHint ?? providerType ?? "",
        currentStep,
        via,
      ).link
    : ORG_DOCS_URL[organizationType];
  const resolvedFooterConfig: WizardFooterConfig = footerConfig;
  const modalTitle = getProviderWizardModalTitle(mode);

  return {
    currentStep,
    docsLink,
    footerConfig,
    handleClose,
    handleDialogOpenChange,
    handleTestSuccess,
    isDirectCredentialsEntry,
    isOrgDirectEntry,
    isProviderFlow,
    mode,
    modalTitle,
    openOrganizationsFlow,
    organizationType,
    orgCurrentStep,
    orgSetupPhase,
    providerTypeHint,
    resolvedFooterConfig,
    setCurrentStep,
    setFooterConfig,
    setOrgCurrentStep,
    setOrgSetupPhase,
    setProviderTypeHint,
    backToProviderFlow,
    wizardVariant,
  };
}