edr-platform/apps/edr-freight-web/portal/src/components/onboarding/OnboardingWizardDialog.tsx

import { Modal, ScrollArea, Stack, Text } from "@mantine/core";
import { useMutation, useQueryClient } from "@tanstack/react-query";
import { useCallback, useState } from "react";

import useAuth from "@/hooks/useAuth";
import { api } from "@/services/api";
import type {
  CreateCompanyPayload,
  ProfileTypeValue,
} from "@/services/companies.service";
import { companiesService } from "@/services/companies.service";
import type { UpdateProfilePayload } from "@/types/profile";
import { extractApiError } from "@/utils/result";
import CompanyProfileForm from "@/pages/accounts/CompanyProfileForm";
import ForwarderForm from "@/pages/accounts/ForwarderForm";
import { FREIGHT_FORWARDER } from "@/pages/settings/companyRoles";
import OnboardingRoleSelect from "@/pages/settings/OnboardingRoleSelect";

/** Form steps shared by CompanyProfileForm and ForwarderForm. */
type FormStep = "company" | "personnel" | "poa" | "documents" | "confirm";
const FORM_STEPS: FormStep[] = [
  "company",
  "personnel",
  "poa",
  "documents",
  "confirm",
];

interface OnboardingWizardDialogProps {
  opened: boolean;
  /** Dismiss the dialog (user clicked the close icon). */
  onClose: () => void;
}

/** Map the chosen operational roles to the company type they belong to. */
function companyTypeForRoles(roles: string[]): string {
  return roles.includes(FREIGHT_FORWARDER.type) ? "forwarder" : "customer";
}

/** Document upload setting code per company type. */
function documentSettingCode(companyType: string): string {
  return companyType === "forwarder"
    ? "company_onboarding_documents_forwarder"
    : "company_onboarding_documents_customer";
}

/**
 * First-run onboarding wizard with a "draft-first" flow: picking the role(s)
 * immediately creates a draft company + profile on the backend, so every
 * subsequent step saves its data incrementally (PATCH /profile, /onboarding-step)
 * against existing rows. The final step uploads documents and marks onboarding
 * complete. Dismissable — the gate keeps it reachable until finished.
 */
export default function OnboardingWizardDialog({
  opened,
  onClose,
}: OnboardingWizardDialogProps) {
  const queryClient = useQueryClient();
  const { user, company, onboardingStep } = useAuth();

  const existingProfiles = company?.company?.companyProfiles ?? [];
  const companyAlreadyStarted = Boolean(company?.company?.id);

  // Resume position from the backend-persisted step.
  const resumeFormStep: FormStep = FORM_STEPS.includes(onboardingStep as FormStep)
    ? (onboardingStep as FormStep)
    : "company";

  // If a draft already exists, resume straight into the form with its roles
  // pre-selected; otherwise start at role selection.
  const [phase, setPhase] = useState<"role" | "form">(
    companyAlreadyStarted ? "form" : "role",
  );
  const [roles, setRoles] = useState<string[]>(
    existingProfiles.map((p) => p.type),
  );
  const [documentFiles, setDocumentFiles] = useState<
    Record<string, File | File[] | null>
  >({});
  const [startError, setStartError] = useState<string | null>(null);

  const refreshInfo = useCallback(
    () =>
      queryClient.invalidateQueries({
        queryKey: api.companies.getInfo.queryKey(),
      }),
    [queryClient],
  );

  // Begin onboarding: create the draft company + profile + role(s).
  const startMutation = useMutation({
    mutationFn: (vars: { companyType: string; roles: ProfileTypeValue[] }) =>
      api.companies.startOnboarding.call(vars),
    onSuccess: async () => {
      await refreshInfo();
      setPhase("form");
    },
    onError: (err) => setStartError(extractApiError(err).message),
  });

  // Finalize: upload any documents, then mark onboarding complete.
  const finishMutation = useMutation({
    mutationFn: async () => {
      const companyId = company?.company?.id;
      const hasFiles = Object.values(documentFiles).some(
        (f) => f !== null && (Array.isArray(f) ? f.length > 0 : true),
      );
      if (companyId && hasFiles) {
        await companiesService.uploadDocuments(companyId, documentFiles);
      }
      return api.companies.completeOnboarding.call();
    },
    onSuccess: refreshInfo,
  });

  // Persist the resume step to the backend (best-effort, fire-and-forget).
  const persistStep = useCallback((step: string) => {
    api.companies.setOnboardingStep.call({ step }).catch(() => {});
  }, []);

  const handleRolesContinue = useCallback(() => {
    setStartError(null);
    startMutation.mutate({
      companyType: companyTypeForRoles(roles),
      roles: roles as ProfileTypeValue[],
    });
  }, [roles, startMutation]);

  const handleBackToRoles = useCallback(() => {
    setPhase("role");
    persistStep("role");
  }, [persistStep]);

  // Save the current step's fields to the draft (PATCH /profile). Returns false
  // to keep the form on the current step when the save fails.
  const saveStep = useCallback(
    async (data: Partial<UpdateProfilePayload>): Promise<boolean> => {
      try {
        await api.companies.updateProfile.call(data as UpdateProfilePayload);
        return true;
      } catch {
        return false;
      }
    },
    [],
  );

  // Final confirm step → finalize onboarding (no company create; it already
  // exists as a draft that's been filled in step-by-step).
  const handleSubmit = useCallback(
    (_payload: CreateCompanyPayload) => {
      finishMutation.mutate();
    },
    [finishMutation],
  );

  if (!user) return null;

  const isForwarder = roles.includes(FREIGHT_FORWARDER.type);
  // Importer+Exporter (or either alone) is a valid customer selection.
  const rolesValid = roles.length > 0;
  const companyType = companyTypeForRoles(roles);

  return (
    <Modal
      opened={opened}
      onClose={onClose}
      withCloseButton
      closeOnClickOutside={false}
      closeOnEscape
      size={1040}
      radius="lg"
      padding="xl"
      centered
      keepMounted
      scrollAreaComponent={ScrollArea.Autosize}
      overlayProps={{ backgroundOpacity: 0.55, blur: 4 }}
      title={
        <Stack gap={2}>
          <Text fz={20} fw={800} c="edr-text" className="tracking-tight">
            Complete your onboarding
          </Text>
          <Text size="sm" c="edr-muted">
            {phase === "role"
              ? "Tell us what your company does to get started."
              : "Set up your company profile to finish."}
          </Text>
        </Stack>
      }
    >
      {phase === "role" ? (
        <Stack gap="lg">
          <OnboardingRoleSelect value={roles} onChange={setRoles} />
          {startError && (
            <Text size="sm" c="red">
              {startError}
            </Text>
          )}
          <RoleContinueBar
            disabled={!rolesValid}
            loading={startMutation.isPending}
            onClick={handleRolesContinue}
          />
        </Stack>
      ) : isForwarder ? (
        <ForwarderForm
          documentSettingCode={documentSettingCode(companyType)}
          documentFiles={documentFiles}
          onDocumentFilesChange={setDocumentFiles}
          user={user}
          onSubmit={handleSubmit}
          isPending={finishMutation.isPending}
          onBack={handleBackToRoles}
          initialStep={resumeFormStep}
          onStepChange={persistStep}
          onSaveStep={saveStep}
        />
      ) : (
        <CompanyProfileForm
          documentSettingCode={documentSettingCode(companyType)}
          documentFiles={documentFiles}
          onDocumentFilesChange={setDocumentFiles}
          user={user}
          onSubmit={handleSubmit}
          isPending={finishMutation.isPending}
          onBack={handleBackToRoles}
          initialStep={resumeFormStep}
          onStepChange={persistStep}
          onSaveStep={saveStep}
        />
      )}
    </Modal>
  );
}

function RoleContinueBar({
  disabled,
  loading,
  onClick,
}: {
  disabled: boolean;
  loading?: boolean;
  onClick: () => void;
}) {
  return (
    <button
      type="button"
      disabled={disabled || loading}
      onClick={onClick}
      className="ml-auto rounded-lg bg-[var(--mantine-color-edr-green-6)] px-5 py-2.5 text-sm font-semibold text-white transition-opacity hover:opacity-90 disabled:cursor-not-allowed disabled:opacity-50"
    >
      {loading ? "Setting up…" : "Continue"}
    </button>
  );
}