wulf-pulse/components/settings/two-factor-setup.tsx

"use client";

import { useState } from "react";
import { Loader2, ShieldCheck, ShieldOff } from "lucide-react";
import { toast } from "sonner";
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogFooter,
  DialogHeader,
  DialogTitle,
} from "@/components/ui/dialog";
import { authClient, useSession } from "@/lib/auth-client";

export function TwoFactorSetup() {
  const { data: session } = useSession();
  const [isLoading, setIsLoading] = useState(false);
  const [showEnableDialog, setShowEnableDialog] = useState(false);
  const [showDisableDialog, setShowDisableDialog] = useState(false);
  const [password, setPassword] = useState("");
  const [totpUri, setTotpUri] = useState("");
  const [verificationCode, setVerificationCode] = useState("");
  const [backupCodes, setBackupCodes] = useState<string[]>([]);
  const [step, setStep] = useState<"password" | "qr" | "verify" | "backup">("password");

  const is2FAEnabled = (session?.user as { twoFactorEnabled?: boolean })?.twoFactorEnabled;

  async function handleEnable() {
    if (step === "password") {
      setIsLoading(true);
      try {
        const result = await authClient.twoFactor.enable({ password });
        if (result.error) {
          toast.error(result.error.message || "Failed to enable 2FA");
          return;
        }
        setTotpUri(result.data?.totpURI || "");
        setBackupCodes(result.data?.backupCodes || []);
        setStep("qr");
      } catch (error) {
        toast.error("An error occurred");
      } finally {
        setIsLoading(false);
      }
    } else if (step === "verify") {
      setIsLoading(true);
      try {
        const result = await authClient.twoFactor.verifyTotp({
          code: verificationCode,
        });
        if (result.error) {
          toast.error(result.error.message || "Invalid code");
          return;
        }
        setStep("backup");
        toast.success("Two-factor authentication enabled!");
      } catch (error) {
        toast.error("An error occurred");
      } finally {
        setIsLoading(false);
      }
    }
  }

  async function handleDisable() {
    setIsLoading(true);
    try {
      const result = await authClient.twoFactor.disable({ password });
      if (result.error) {
        toast.error(result.error.message || "Failed to disable 2FA");
        return;
      }
      toast.success("Two-factor authentication disabled");
      setShowDisableDialog(false);
      setPassword("");
    } catch (error) {
      toast.error("An error occurred");
    } finally {
      setIsLoading(false);
    }
  }

  function resetDialog() {
    setStep("password");
    setPassword("");
    setTotpUri("");
    setVerificationCode("");
    setBackupCodes([]);
  }

  return (
    <div>
      {is2FAEnabled ? (
        <div className="flex items-center justify-between">
          <div className="flex items-center gap-3">
            <div className="h-10 w-10 rounded-full bg-green-100 dark:bg-green-900/30 flex items-center justify-center">
              <ShieldCheck className="h-5 w-5 text-green-600 dark:text-green-400" />
            </div>
            <div>
              <p className="font-medium">Two-factor authentication is enabled</p>
              <p className="text-sm text-muted-foreground">
                Your account is protected with an authenticator app
              </p>
            </div>
          </div>
          <Button variant="outline" onClick={() => setShowDisableDialog(true)}>
            Disable
          </Button>
        </div>
      ) : (
        <div className="flex items-center justify-between">
          <div className="flex items-center gap-3">
            <div className="h-10 w-10 rounded-full bg-muted flex items-center justify-center">
              <ShieldOff className="h-5 w-5 text-muted-foreground" />
            </div>
            <div>
              <p className="font-medium">Two-factor authentication is disabled</p>
              <p className="text-sm text-muted-foreground">
                Add an extra layer of security to your account
              </p>
            </div>
          </div>
          <Button onClick={() => setShowEnableDialog(true)}>Enable</Button>
        </div>
      )}

      {/* Enable 2FA Dialog */}
      <Dialog open={showEnableDialog} onOpenChange={(open) => {
        setShowEnableDialog(open);
        if (!open) resetDialog();
      }}>
        <DialogContent>
          <DialogHeader>
            <DialogTitle>
              {step === "password" && "Enable Two-Factor Authentication"}
              {step === "qr" && "Scan QR Code"}
              {step === "verify" && "Verify Code"}
              {step === "backup" && "Save Backup Codes"}
            </DialogTitle>
            <DialogDescription>
              {step === "password" && "Enter your password to continue"}
              {step === "qr" && "Scan this QR code with your authenticator app"}
              {step === "verify" && "Enter the 6-digit code from your authenticator app"}
              {step === "backup" && "Save these backup codes in a safe place"}
            </DialogDescription>
          </DialogHeader>

          {step === "password" && (
            <div className="space-y-4">
              <div className="space-y-2">
                <Label htmlFor="password">Password</Label>
                <Input
                  id="password"
                  type="password"
                  value={password}
                  onChange={(e) => setPassword(e.target.value)}
                  disabled={isLoading}
                />
              </div>
            </div>
          )}

          {step === "qr" && (
            <div className="space-y-4">
              <div className="flex justify-center p-4 bg-white rounded-lg">
                {/* QR Code would be rendered here - using a placeholder */}
                <div className="text-center">
                  <p className="text-sm text-muted-foreground mb-2">
                    Scan with your authenticator app or enter manually:
                  </p>
                  <code className="text-xs break-all">{totpUri}</code>
                </div>
              </div>
              <Button className="w-full" onClick={() => setStep("verify")}>
                Continue
              </Button>
            </div>
          )}

          {step === "verify" && (
            <div className="space-y-4">
              <div className="space-y-2">
                <Label htmlFor="code">Verification Code</Label>
                <Input
                  id="code"
                  value={verificationCode}
                  onChange={(e) => setVerificationCode(e.target.value)}
                  placeholder="000000"
                  maxLength={6}
                  className="text-center text-2xl tracking-widest"
                  disabled={isLoading}
                />
              </div>
            </div>
          )}

          {step === "backup" && (
            <div className="space-y-4">
              <div className="bg-muted p-4 rounded-lg">
                <div className="grid grid-cols-2 gap-2">
                  {backupCodes.map((code, i) => (
                    <code key={i} className="text-sm font-mono">
                      {code}
                    </code>
                  ))}
                </div>
              </div>
              <p className="text-sm text-muted-foreground">
                Each code can only be used once. Store them securely.
              </p>
            </div>
          )}

          <DialogFooter>
            {step === "backup" ? (
              <Button onClick={() => {
                setShowEnableDialog(false);
                resetDialog();
              }}>
                Done
              </Button>
            ) : step !== "qr" && (
              <Button onClick={handleEnable} disabled={isLoading}>
                {isLoading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
                {step === "password" ? "Continue" : "Verify"}
              </Button>
            )}
          </DialogFooter>
        </DialogContent>
      </Dialog>

      {/* Disable 2FA Dialog */}
      <Dialog open={showDisableDialog} onOpenChange={setShowDisableDialog}>
        <DialogContent>
          <DialogHeader>
            <DialogTitle>Disable Two-Factor Authentication</DialogTitle>
            <DialogDescription>
              Enter your password to disable two-factor authentication
            </DialogDescription>
          </DialogHeader>
          <div className="space-y-4">
            <div className="space-y-2">
              <Label htmlFor="disable-password">Password</Label>
              <Input
                id="disable-password"
                type="password"
                value={password}
                onChange={(e) => setPassword(e.target.value)}
                disabled={isLoading}
              />
            </div>
          </div>
          <DialogFooter>
            <Button variant="outline" onClick={() => setShowDisableDialog(false)}>
              Cancel
            </Button>
            <Button variant="destructive" onClick={handleDisable} disabled={isLoading}>
              {isLoading && <Loader2 className="mr-2 h-4 w-4 animate-spin" />}
              Disable
            </Button>
          </DialogFooter>
        </DialogContent>
      </Dialog>
    </div>
  );
}