CedarCopilot/cedar-mail

import { twoFactor } from '@/modules/auth/utils/auth-client';
import { useCallback, useMemo, useState } from 'react';
import { toast } from 'sonner';

export type TwoFactorEnrollment = { totpURI: string; backupCodes: string[] };

/**
 * The enrollment and challenge calls, shared by the forced `/two-factor` interstitial and
 * the security settings screen.
 *
 * Shared deliberately: both screens set up the same factor, and a second copy of "call
 * enable, hold the backup codes, verify once" is exactly the kind of duplication that
 * quietly diverges — one screen growing a fix the other never gets.
 */
export function useTwoFactor() {
  const [enrollment, setEnrollment] = useState<TwoFactorEnrollment | null>(null);
  const [busy, setBusy] = useState(false);

  /**
   * Starts enrollment and returns the secret plus one-time recovery codes. No password
   * argument: Cedar has no passwords, and the server registers the plugin with
   * `allowPasswordless` so it does not ask for one.
   */
  const begin = useCallback(async () => {
    setBusy(true);
    try {
      const { data, error } = await twoFactor.enable({});
      if (error || !data) {
        toast.error(error?.message ?? 'Could not start two-factor setup.');
        return false;
      }
      setEnrollment({ totpURI: data.totpURI, backupCodes: data.backupCodes });
      return true;
    } finally {
      setBusy(false);
    }
  }, []);

  /**
   * Verifies a code. The same call both completes a first enrollment and clears a later
   * challenge — better-auth decides which from the account's state, not from a flag here.
   */
  const verify = useCallback(async (code: string, options?: { useBackupCode?: boolean }) => {
    const trimmed = code.trim();
    if (!trimmed) return false;

    setBusy(true);
    try {
      const { error } = options?.useBackupCode
        ? await twoFactor.verifyBackupCode({ code: trimmed })
        : await twoFactor.verifyTotp({ code: trimmed });

      if (error) {
        toast.error(
          options?.useBackupCode
            ? 'That recovery code was not accepted.'
            : 'That code was not accepted.',
        );
        return false;
      }
      return true;
    } finally {
      setBusy(false);
    }
  }, []);

  /** Issues a fresh set of recovery codes, invalidating the previous set. */
  const regenerateBackupCodes = useCallback(async () => {
    setBusy(true);
    try {
      const { data, error } = await twoFactor.generateBackupCodes({});
      if (error || !data) {
        toast.error(error?.message ?? 'Could not generate new recovery codes.');
        return null;
      }
      return data.backupCodes;
    } finally {
      setBusy(false);
    }
  }, []);

  /**
   * Turns the factor off. The server rejects this outright for users the policy requires it
   * of, so a refusal here is expected rather than exceptional.
   */
  const disable = useCallback(async () => {
    setBusy(true);
    try {
      const { error } = await twoFactor.disable({});
      if (error) {
        toast.error(error.message ?? 'Could not turn off two-factor authentication.');
        return false;
      }
      return true;
    } finally {
      setBusy(false);
    }
  }, []);

  const manualEntryKey=[redacted] => extractSecret(enrollment?.totpURI), [enrollment]);

  const clearEnrollment = useCallback(() => setEnrollment(null), []);

  return {
    enrollment,
    manualEntryKey,
    busy,
    begin,
    verify,
    regenerateBackupCodes,
    disable,
    clearEnrollment,
  };
}

/**
 * Pulls the base32 secret out of an `otpauth://` URI for manual entry. Regex rather than
 * `new URL`, because URL's searchParams parsing is not dependable for non-http schemes
 * across browsers.
 */
export function extractSecret(totpURI: string | undefined): string | null {
  if (!totpURI) return null;
  return /[?&]secret=[redacted] ?? null;
}