Nexus/apps/web/src/components/security/MfaSetup.tsx

"use client";

import { useState, useEffect } from "react";
import QRCode from "qrcode";
import { trpc } from "~/lib/trpc/client.js";

type SetupStep = "idle" | "show-secret" | "verify" | "show-backup-codes" | "done";

export function MfaSetup() {
  const [step, setStep] = useState<SetupStep>("idle");
  const [secret, setSecret] = useState("");
  const [uri, setUri] = useState("");
  const [qrDataUrl, setQrDataUrl] = useState("");
  const [token, setToken] = useState("");
  const [backupCodes, setBackupCodes] = useState<string[] | null>(null);
  const [error, setError] = useState<string | null>(null);
  const [success, setSuccess] = useState<string | null>(null);

  useEffect(() => {
    if (!uri) return;
    let cancelled = false;
    QRCode.toDataURL(uri, { width: 200, margin: 2 })
      .then((dataUrl) => {
        if (!cancelled) setQrDataUrl(dataUrl);
      })
      .catch(() => {
        /* ignore — manual key is shown as fallback */
      });
    return () => {
      cancelled = true;
    };
  }, [uri]);

  const { data: mfaStatus, refetch } = trpc.user.getMfaStatus.useQuery();
  const generateMutation = trpc.user.generateTotpSecret.useMutation();
  const verifyMutation = trpc.user.verifyAndEnableTotp.useMutation();
  const regenerateBackupCodesMutation = trpc.user.regenerateBackupCodes.useMutation();

  async function handleGenerate() {
    setError(null);
    try {
      const result = await generateMutation.mutateAsync();
      setSecret(result.secret);
      setUri(result.uri);
      setStep("show-secret");
    } catch (err) {
      setError(err instanceof Error ? err.message : "Failed to generate TOTP secret");
    }
  }

  async function handleVerify() {
    setError(null);
    try {
      const result = await verifyMutation.mutateAsync({ token });
      setBackupCodes(result.backupCodes ?? null);
      setStep("show-backup-codes");
      setSecret("");
      setUri("");
      setToken("");
      await refetch();
    } catch (err) {
      setError(err instanceof Error ? err.message : "Verification failed");
    }
  }

  async function handleRegenerateBackupCodes() {
    setError(null);
    try {
      const result = await regenerateBackupCodesMutation.mutateAsync();
      setBackupCodes(result.codes);
      setStep("show-backup-codes");
      await refetch();
    } catch (err) {
      setError(err instanceof Error ? err.message : "Could not regenerate backup codes");
    }
  }

  function handleFinishBackupCodes() {
    setBackupCodes(null);
    setStep("done");
    setSuccess("MFA is active. Keep your backup codes in a safe place.");
  }

  function copyBackupCodes() {
    if (!backupCodes) return;
    void navigator.clipboard.writeText(backupCodes.join("\n"));
  }

  function downloadBackupCodes() {
    if (!backupCodes) return;
    const blob = new Blob(
      [
        `Nexus MFA Backup Codes\nGenerated: ${new Date().toISOString()}\n\nEach code works exactly once. Keep this file somewhere safe.\n\n${backupCodes.join("\n")}\n`,
      ],
      { type: "text/plain" },
    );
    const url = URL.createObjectURL(blob);
    const a = document.createElement("a");
    a.href = url;
    a.download = "nexus-backup-codes.txt";
    a.click();
    URL.revokeObjectURL(url);
  }

  if (mfaStatus?.totpEnabled && step !== "done" && step !== "show-backup-codes") {
    const remaining = mfaStatus.backupCodesRemaining ?? 0;
    const lowCodes = remaining <= 3;
    return (
      <div className="space-y-4">
        <div className="rounded-xl border border-green-200 dark:border-green-800 bg-green-50 dark:bg-green-900/20 p-6">
          <div className="flex items-center gap-3">
            <div className="flex h-10 w-10 items-center justify-center rounded-full bg-green-100 dark:bg-green-900/40">
              <svg
                className="h-5 w-5 text-green-600 dark:text-green-400"
                fill="none"
                stroke="currentColor"
                viewBox="0 0 24 24"
              >
                <path
                  strokeLinecap="round"
                  strokeLinejoin="round"
                  strokeWidth={2}
                  d="M9 12l2 2 4-4m5.618-4.016A11.955 11.955 0 0112 2.944a11.955 11.955 0 01-8.618 3.04A12.02 12.02 0 003 9c0 5.591 3.824 10.29 9 11.622 5.176-1.332 9-6.03 9-11.622 0-1.042-.133-2.052-.382-3.016z"
                />
              </svg>
            </div>
            <div>
              <h3 className="text-sm font-semibold text-green-800 dark:text-green-300">
                MFA Enabled
              </h3>
              <p className="text-sm text-green-700 dark:text-green-400">
                Two-factor authentication is active on your account.
              </p>
            </div>
          </div>
        </div>

        <div
          className={`rounded-xl border p-6 ${
            lowCodes
              ? "border-amber-200 dark:border-amber-800 bg-amber-50 dark:bg-amber-900/20"
              : "border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-900"
          }`}
        >
          <div className="flex items-start justify-between gap-4">
            <div>
              <h3 className="text-sm font-semibold text-gray-900 dark:text-gray-100">
                Backup codes
              </h3>
              <p className="mt-1 text-sm text-gray-600 dark:text-gray-400">
                {remaining === 0
                  ? "You have no backup codes left. Generate a new set to avoid being locked out if you lose your device."
                  : `You have ${remaining} backup code${remaining === 1 ? "" : "s"} remaining.`}{" "}
                {lowCodes && remaining > 0 && <span className="font-medium">Regenerate soon.</span>}
              </p>
            </div>
            <button
              type="button"
              onClick={handleRegenerateBackupCodes}
              disabled={regenerateBackupCodesMutation.isPending}
              className="shrink-0 inline-flex items-center gap-2 rounded-lg border border-gray-300 dark:border-gray-600 bg-white dark:bg-gray-800 px-3 py-2 text-sm font-medium text-gray-700 dark:text-gray-200 hover:bg-gray-50 dark:hover:bg-gray-700 disabled:opacity-50"
            >
              {regenerateBackupCodesMutation.isPending ? "Generating…" : "Regenerate codes"}
            </button>
          </div>
          {error && (
            <div className="mt-3 rounded-lg bg-red-50 dark:bg-red-900/20 border border-red-200 dark:border-red-700 px-4 py-2 text-sm text-red-700 dark:text-red-400">
              {error}
            </div>
          )}
        </div>
      </div>
    );
  }

  return (
    <div className="space-y-4">
      {error && (
        <div className="rounded-lg bg-red-50 dark:bg-red-900/20 border border-red-200 dark:border-red-700 px-4 py-3 text-sm text-red-700 dark:text-red-400">
          {error}
        </div>
      )}
      {success && (
        <div className="rounded-lg bg-green-50 dark:bg-green-900/20 border border-green-200 dark:border-green-700 px-4 py-3 text-sm text-green-700 dark:text-green-400">
          {success}
        </div>
      )}

      {step === "idle" && (
        <div className="rounded-xl border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-900 p-6">
          <div className="flex items-start gap-4">
            <div className="flex h-10 w-10 shrink-0 items-center justify-center rounded-full bg-amber-100 dark:bg-amber-900/40">
              <svg
                className="h-5 w-5 text-amber-600 dark:text-amber-400"
                fill="none"
                stroke="currentColor"
                viewBox="0 0 24 24"
              >
                <path
                  strokeLinecap="round"
                  strokeLinejoin="round"
                  strokeWidth={2}
                  d="M12 15v2m-6 4h12a2 2 0 002-2v-6a2 2 0 00-2-2H6a2 2 0 00-2 2v6a2 2 0 002 2zm10-10V7a4 4 0 00-8 0v4h8z"
                />
              </svg>
            </div>
            <div>
              <h3 className="text-sm font-semibold text-gray-900 dark:text-gray-100">
                Two-Factor Authentication (TOTP)
              </h3>
              <p className="mt-1 text-sm text-gray-600 dark:text-gray-400">
                Add an extra layer of security by requiring a code from your authenticator app when
                signing in.
              </p>
              <button
                type="button"
                onClick={handleGenerate}
                disabled={generateMutation.isPending}
                className="mt-4 inline-flex items-center gap-2 rounded-lg bg-brand-600 px-4 py-2 text-sm font-medium text-white shadow-sm hover:bg-brand-700 transition-colors disabled:opacity-50"
              >
                {generateMutation.isPending ? "Generating..." : "Set up MFA"}
              </button>
            </div>
          </div>
        </div>
      )}

      {step === "show-secret" && (
        <div className="rounded-xl border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-900 p-6 space-y-5">
          <h3 className="text-sm font-semibold text-gray-900 dark:text-gray-100">
            Step 1: Scan the QR code
          </h3>
          <p className="text-sm text-gray-600 dark:text-gray-400">
            Scan this QR code with your authenticator app (Google Authenticator, Authy, 1Password,
            etc.).
          </p>

          {/* QR Code — rendered locally, no external service */}
          <div className="flex justify-center">
            <div className="rounded-lg border border-gray-200 dark:border-gray-700 bg-white p-3">
              {qrDataUrl ? (
                <img
                  src={qrDataUrl}
                  alt="TOTP QR Code"
                  width={200}
                  height={200}
                  className="rounded"
                />
              ) : (
                <div className="h-[200px] w-[200px] flex items-center justify-center text-xs text-gray-400">
                  Generating…
                </div>
              )}
            </div>
          </div>

          <div>
            <p className="text-xs font-medium text-gray-500 dark:text-gray-400 mb-1">
              Or enter this key manually:
            </p>
            <code className="block rounded-lg bg-gray-100 dark:bg-gray-800 px-4 py-2 text-sm font-mono text-gray-900 dark:text-gray-100 break-all select-all">
              {secret}
            </code>
          </div>

          <button
            type="button"
            onClick={() => {
              setStep("verify");
              setError(null);
            }}
            className="inline-flex items-center gap-2 rounded-lg bg-brand-600 px-4 py-2 text-sm font-medium text-white shadow-sm hover:bg-brand-700 transition-colors"
          >
            Continue
          </button>
        </div>
      )}

      {step === "verify" && (
        <div className="rounded-xl border border-gray-200 dark:border-gray-700 bg-white dark:bg-gray-900 p-6 space-y-5">
          <h3 className="text-sm font-semibold text-gray-900 dark:text-gray-100">
            Step 2: Verify your code
          </h3>
          <p className="text-sm text-gray-600 dark:text-gray-400">
            Enter the 6-digit code from your authenticator app to confirm setup.
          </p>

          <div>
            <label
              htmlFor="mfa-verify-token"
              className="block text-sm font-medium text-gray-700 dark:text-gray-300 mb-1"
            >
              Verification Code
            </label>
            <input
              id="mfa-verify-token"
              type="text"
              inputMode="numeric"
              autoComplete="one-time-code"
              maxLength={6}
              pattern="[0-9]{6}"
              value={token}
              onChange={(e) => setToken(e.target.value.replace(/\D/g, "").slice(0, 6))}
              className="w-48 rounded-lg border border-gray-300 dark:border-gray-600 bg-white dark:bg-gray-800 px-4 py-2 text-center text-xl font-mono tracking-[0.3em] text-gray-900 dark:text-gray-100 focus:outline-none focus:ring-2 focus:ring-brand-400"
              placeholder="000000"
              autoFocus
            />
          </div>

          <div className="flex items-center gap-3">
            <button
              type="button"
              onClick={handleVerify}
              disabled={token.length !== 6 || verifyMutation.isPending}
              className="inline-flex items-center gap-2 rounded-lg bg-brand-600 px-4 py-2 text-sm font-medium text-white shadow-sm hover:bg-brand-700 transition-colors disabled:opacity-50"
            >
              {verifyMutation.isPending ? "Verifying..." : "Enable MFA"}
            </button>
            <button
              type="button"
              onClick={() => {
                setStep("show-secret");
                setToken("");
                setError(null);
              }}
              className="text-sm text-gray-500 hover:text-gray-700 dark:text-gray-400 dark:hover:text-gray-200"
            >
              Back
            </button>
          </div>
        </div>
      )}

      {step === "show-backup-codes" && backupCodes && (
        <div className="rounded-xl border border-amber-200 dark:border-amber-800 bg-amber-50 dark:bg-amber-900/20 p-6 space-y-4">
          <div>
            <h3 className="text-sm font-semibold text-amber-900 dark:text-amber-200">
              Save your backup codes
            </h3>
            <p className="mt-1 text-sm text-amber-800 dark:text-amber-300">
              Each code works exactly once. Store them in a password manager or print them. You will
              not see them again — regenerating invalidates the whole set.
            </p>
          </div>
          <div className="grid grid-cols-2 gap-2 rounded-lg bg-white dark:bg-gray-900 p-4 font-mono text-sm">
            {backupCodes.map((code) => (
              <code
                key={code}
                className="rounded bg-gray-100 dark:bg-gray-800 px-3 py-2 text-center tracking-wider select-all"
              >
                {code}
              </code>
            ))}
          </div>
          <div className="flex flex-wrap items-center gap-2">
            <button
              type="button"
              onClick={copyBackupCodes}
              className="inline-flex items-center gap-2 rounded-lg border border-gray-300 dark:border-gray-600 bg-white dark:bg-gray-800 px-3 py-2 text-sm font-medium text-gray-700 dark:text-gray-200 hover:bg-gray-50 dark:hover:bg-gray-700"
            >
              Copy all
            </button>
            <button
              type="button"
              onClick={downloadBackupCodes}
              className="inline-flex items-center gap-2 rounded-lg border border-gray-300 dark:border-gray-600 bg-white dark:bg-gray-800 px-3 py-2 text-sm font-medium text-gray-700 dark:text-gray-200 hover:bg-gray-50 dark:hover:bg-gray-700"
            >
              Download .txt
            </button>
            <button
              type="button"
              onClick={handleFinishBackupCodes}
              className="ml-auto inline-flex items-center gap-2 rounded-lg bg-brand-600 px-4 py-2 text-sm font-medium text-white shadow-sm hover:bg-brand-700"
            >
              I've saved them
            </button>
          </div>
        </div>
      )}
    </div>
  );
}