myeasycms-v2/packages/features/auth/src/components/multi-factor-challenge-container.tsx

'use client';

import { useCallback, useEffect, useEffectEvent, useState } from 'react';

import { useRouter } from 'next/navigation';

import type { Factor } from '@supabase/supabase-js';

import { zodResolver } from '@hookform/resolvers/zod';
import { useMutation } from '@tanstack/react-query';
import { Fingerprint, KeyRound, TriangleAlert } from 'lucide-react';
import { useForm } from 'react-hook-form';
import * as z from 'zod';

import { useFetchAuthFactors } from '@kit/supabase/hooks/use-fetch-mfa-factors';
import { useSignOut } from '@kit/supabase/hooks/use-sign-out';
import { useSupabase } from '@kit/supabase/hooks/use-supabase';
import { Alert, AlertDescription, AlertTitle } from '@kit/ui/alert';
import { Button } from '@kit/ui/button';
import {
  Form,
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormMessage,
} from '@kit/ui/form';
import { Heading } from '@kit/ui/heading';
import { If } from '@kit/ui/if';
import {
  InputOTP,
  InputOTPGroup,
  InputOTPSeparator,
  InputOTPSlot,
} from '@kit/ui/input-otp';
import { Spinner } from '@kit/ui/spinner';
import { Trans } from '@kit/ui/trans';

export function MultiFactorChallengeContainer({
  paths,
  userId,
}: React.PropsWithChildren<{
  userId: string;
  paths: {
    redirectPath: string;
  };
}>) {
  const router = useRouter();
  const [selectedFactor, setSelectedFactor] = useState<Factor | null>(null);

  const onSuccess = useCallback(() => {
    router.replace(paths.redirectPath);
  }, [router, paths.redirectPath]);

  if (!selectedFactor) {
    return (
      <FactorsListContainer userId={userId} onSelect={setSelectedFactor} />
    );
  }

  if (selectedFactor.factor_type === 'webauthn') {
    return (
      <WebAuthnChallengeView
        factor={selectedFactor}
        onSuccess={onSuccess}
        onBack={() => setSelectedFactor(null)}
      />
    );
  }

  return <TotpChallengeView factor={selectedFactor} onSuccess={onSuccess} />;
}

function TotpChallengeView({
  factor,
  onSuccess,
}: {
  factor: Factor;
  onSuccess: () => void;
}) {
  const verifyMFAChallenge = useVerifyMFAChallenge({ onSuccess });

  const verificationCodeForm = useForm({
    resolver: zodResolver(
      z.object({
        verificationCode: z.string().min(6).max(6),
      }),
    ),
    defaultValues: {
      verificationCode: '',
    },
  });

  return (
    <Form {...verificationCodeForm}>
      <form
        className={'w-full'}
        onSubmit={verificationCodeForm.handleSubmit(async (data) => {
          await verifyMFAChallenge.mutateAsync({
            factorId: factor.id,
            verificationCode: data.verificationCode,
          });
        })}
      >
        <div className={'flex flex-col items-center gap-y-6'}>
          <div className="flex flex-col items-center gap-y-4">
            <Heading level={5}>
              <Trans i18nKey={'auth.verifyCodeHeading'} />
            </Heading>
          </div>

          <div className={'flex w-full flex-col gap-y-2.5'}>
            <div className={'flex flex-col gap-y-4'}>
              <If condition={verifyMFAChallenge.error}>
                <Alert variant={'destructive'}>
                  <TriangleAlert className={'h-5'} />

                  <AlertTitle>
                    <Trans i18nKey={'account.invalidVerificationCodeHeading'} />
                  </AlertTitle>

                  <AlertDescription>
                    <Trans
                      i18nKey={'account.invalidVerificationCodeDescription'}
                    />
                  </AlertDescription>
                </Alert>
              </If>

              <FormField
                name={'verificationCode'}
                render={({ field }) => {
                  return (
                    <FormItem
                      className={
                        'mx-auto flex flex-col items-center justify-center'
                      }
                    >
                      <FormControl>
                        <InputOTP {...field} maxLength={6} minLength={6}>
                          <InputOTPGroup>
                            <InputOTPSlot index={0} />
                            <InputOTPSlot index={1} />
                            <InputOTPSlot index={2} />
                          </InputOTPGroup>
                          <InputOTPSeparator />
                          <InputOTPGroup>
                            <InputOTPSlot index={3} />
                            <InputOTPSlot index={4} />
                            <InputOTPSlot index={5} />
                          </InputOTPGroup>
                        </InputOTP>
                      </FormControl>

                      <FormDescription className="text-center">
                        <Trans
                          i18nKey={'account.verifyActivationCodeDescription'}
                        />
                      </FormDescription>

                      <FormMessage />
                    </FormItem>
                  );
                }}
              />
            </div>
          </div>

          <Button
            type="submit"
            className="w-full"
            data-test={'submit-mfa-button'}
            disabled={
              verifyMFAChallenge.isPending ||
              verifyMFAChallenge.isSuccess ||
              !verificationCodeForm.formState.isValid
            }
          >
            <If condition={verifyMFAChallenge.isPending}>
              <span className={'animate-in fade-in slide-in-from-bottom-24'}>
                <Trans i18nKey={'account.verifyingCode'} />
              </span>
            </If>

            <If condition={verifyMFAChallenge.isSuccess}>
              <span className={'animate-in fade-in slide-in-from-bottom-24'}>
                <Trans i18nKey={'auth.redirecting'} />
              </span>
            </If>

            <If
              condition={
                !verifyMFAChallenge.isPending && !verifyMFAChallenge.isSuccess
              }
            >
              <Trans i18nKey={'account.submitVerificationCode'} />
            </If>
          </Button>
        </div>
      </form>
    </Form>
  );
}

function WebAuthnChallengeView({
  factor,
  onSuccess,
  onBack,
}: {
  factor: Factor;
  onSuccess: () => void;
  onBack: () => void;
}) {
  const authenticatePasskey = useAuthenticatePasskey({ onSuccess });
  const [autoTriggered, setAutoTriggered] = useState(false);

  // Auto-trigger the passkey challenge on mount
  useEffect(() => {
    if (!autoTriggered) {
      setAutoTriggered(true);
      authenticatePasskey.mutate(factor.id);
    }
  }, [autoTriggered, authenticatePasskey, factor.id]);

  return (
    <div className={'flex flex-col items-center gap-y-6'}>
      <div className="flex flex-col items-center gap-y-4">
        <Heading level={5}>
          <Trans i18nKey={'account.passkeyVerifyHeading'} />
        </Heading>
      </div>

      <If condition={authenticatePasskey.error}>
        <Alert variant={'destructive'}>
          <TriangleAlert className={'h-5'} />

          <AlertTitle>
            <Trans i18nKey={'account.passkeyVerifyErrorHeading'} />
          </AlertTitle>

          <AlertDescription>
            <Trans i18nKey={'account.passkeyVerifyErrorDescription'} />
          </AlertDescription>
        </Alert>
      </If>

      <div
        className={'flex flex-col items-center space-y-2 rounded-lg border p-6'}
      >
        <Fingerprint className={'text-muted-foreground h-12 w-12'} />

        <p className={'text-muted-foreground text-center text-sm'}>
          <If condition={authenticatePasskey.isPending}>
            <Trans i18nKey={'account.passkeyVerifyWaiting'} />
          </If>

          <If condition={authenticatePasskey.isSuccess}>
            <Trans i18nKey={'auth.redirecting'} />
          </If>

          <If
            condition={
              !authenticatePasskey.isPending && !authenticatePasskey.isSuccess
            }
          >
            <Trans i18nKey={'account.passkeyVerifyPrompt'} />
          </If>
        </p>
      </div>

      <div className={'flex w-full gap-2'}>
        <Button
          variant={'outline'}
          className={'flex-1'}
          onClick={onBack}
          disabled={authenticatePasskey.isPending}
        >
          <Trans i18nKey={'common.back'} />
        </Button>

        <Button
          className={'flex-1'}
          onClick={() => authenticatePasskey.mutate(factor.id)}
          disabled={
            authenticatePasskey.isPending || authenticatePasskey.isSuccess
          }
        >
          <Fingerprint className={'h-4'} />
          <Trans i18nKey={'account.passkeyRetry'} />
        </Button>
      </div>
    </div>
  );
}

function useVerifyMFAChallenge({ onSuccess }: { onSuccess: () => void }) {
  const client = useSupabase();
  const mutationKey = ['mfa-verify-challenge'];

  const mutationFn = async (params: {
    factorId: string;
    verificationCode: string;
  }) => {
    const { factorId, verificationCode: code } = params;

    const response = await client.auth.mfa.challengeAndVerify({
      factorId,
      code,
    });

    if (response.error) {
      throw response.error;
    }

    return response.data;
  };

  return useMutation({ mutationKey, mutationFn, onSuccess });
}

function useAuthenticatePasskey({ onSuccess }: { onSuccess: () => void }) {
  const client = useSupabase();
  const mutationKey = ['mfa-webauthn-authenticate'];

  const mutationFn = async (factorId: string) => {
    const { data, error } = await client.auth.mfa.webauthn.authenticate({
      factorId,
    });

    if (error) {
      throw error;
    }

    return data;
  };

  return useMutation({ mutationKey, mutationFn, onSuccess });
}

function FactorsListContainer({
  onSelect,
  userId,
}: React.PropsWithChildren<{
  userId: string;
  onSelect: (factor: Factor) => void;
}>) {
  const signOut = useSignOut();
  const { data: factors, isLoading, error } = useFetchAuthFactors(userId);

  const isSuccess = factors && !isLoading && !error;

  const signOutFn = useEffectEvent(() => {
    void signOut.mutateAsync();
  });

  useEffect(() => {
    // If there is an error, sign out
    if (error) {
      void signOutFn();
    }
  }, [error]);

  useEffect(() => {
    if (!isSuccess) return;

    const allVerified = [
      ...(factors.totp ?? []),
      ...((factors as Record<string, Factor[]>).webauthn ?? []),
    ];

    // If there is only one factor, select it automatically
    if (allVerified.length === 1 && allVerified[0]) {
      onSelect(allVerified[0]);
    }
  });

  if (isLoading) {
    return (
      <div className={'flex flex-col items-center space-y-4 py-8'}>
        <Spinner />

        <div className={'text-sm'}>
          <Trans i18nKey={'account.loadingFactors'} />
        </div>
      </div>
    );
  }

  if (error) {
    return (
      <div className={'w-full'}>
        <Alert variant={'destructive'}>
          <TriangleAlert className={'h-4'} />

          <AlertTitle>
            <Trans i18nKey={'account.factorsListError'} />
          </AlertTitle>

          <AlertDescription>
            <Trans i18nKey={'account.factorsListErrorDescription'} />
          </AlertDescription>
        </Alert>
      </div>
    );
  }

  const totpFactors = factors?.totp ?? [];
  const webauthnFactors =
    (factors as Record<string, Factor[]> | undefined)?.webauthn ?? [];

  return (
    <div className={'animate-in fade-in flex flex-col space-y-4 duration-500'}>
      <div>
        <span className={'font-medium'}>
          <Trans i18nKey={'account.selectFactor'} />
        </span>
      </div>

      <div className={'flex flex-col space-y-2'}>
        {totpFactors.map((factor) => (
          <div key={factor.id}>
            <Button
              variant={'outline'}
              className={'w-full justify-start gap-2'}
              onClick={() => onSelect(factor)}
            >
              <KeyRound className={'h-4 w-4'} />
              {factor.friendly_name}
            </Button>
          </div>
        ))}

        {webauthnFactors.map((factor) => (
          <div key={factor.id}>
            <Button
              variant={'outline'}
              className={'w-full justify-start gap-2'}
              onClick={() => onSelect(factor)}
            >
              <Fingerprint className={'h-4 w-4'} />
              {factor.friendly_name}
            </Button>
          </div>
        ))}
      </div>
    </div>
  );
}