smoothschedule/frontend/src/pages/MFAVerifyPage.tsx

/**
 * MFA Verification Page
 * Shown when user has MFA enabled and needs to complete verification during login
 */

import React, { useState, useEffect, useRef } from 'react';
import { useTranslation } from 'react-i18next';
import { useNavigate } from 'react-router-dom';
import { sendMFALoginCode, verifyMFALogin } from '../api/mfa';
import { setCookie } from '../utils/cookies';
import { buildSubdomainUrl } from '../utils/domain';
import SmoothScheduleLogo from '../components/SmoothScheduleLogo';
import {
  AlertCircle,
  Loader2,
  Shield,
  Smartphone,
  Key,
  ArrowLeft,
  CheckCircle
} from 'lucide-react';

interface MFAChallenge {
  user_id: number;
  mfa_methods: ('SMS' | 'TOTP' | 'BACKUP')[];
  phone_last_4: string | null;
}

const MFAVerifyPage: React.FC = () => {
  const { t } = useTranslation();
  const navigate = useNavigate();

  const [challenge, setChallenge] = useState<MFAChallenge | null>(null);
  const [selectedMethod, setSelectedMethod] = useState<'SMS' | 'TOTP' | 'BACKUP' | null>(null);
  const [code, setCode] = useState(['', '', '', '', '', '']);
  const [backupCode, setBackupCode] = useState('');
  const [error, setError] = useState('');
  const [loading, setLoading] = useState(false);
  const [smsSent, setSmsSent] = useState(false);
  const [trustDevice, setTrustDevice] = useState(false);

  const inputRefs = useRef<(HTMLInputElement | null)[]>([]);

  useEffect(() => {
    // Get MFA challenge from sessionStorage
    const storedChallenge = sessionStorage.getItem('mfa_challenge');
    if (!storedChallenge) {
      navigate('/login');
      return;
    }

    const parsed = JSON.parse(storedChallenge) as MFAChallenge;
    setChallenge(parsed);

    // Default to TOTP if available, otherwise SMS
    if (parsed.mfa_methods.includes('TOTP')) {
      setSelectedMethod('TOTP');
    } else if (parsed.mfa_methods.includes('SMS')) {
      setSelectedMethod('SMS');
    }
  }, [navigate]);

  const handleSendSMS = async () => {
    if (!challenge) return;

    setLoading(true);
    setError('');

    try {
      await sendMFALoginCode(challenge.user_id, 'SMS');
      setSmsSent(true);
    } catch (err: any) {
      setError(err.response?.data?.error || 'Failed to send SMS code');
    } finally {
      setLoading(false);
    }
  };

  const handleCodeChange = (index: number, value: string) => {
    // Only allow digits
    if (value && !/^\d$/.test(value)) return;

    const newCode = [...code];
    newCode[index] = value;
    setCode(newCode);

    // Auto-focus next input
    if (value && index < 5) {
      inputRefs.current[index + 1]?.focus();
    }
  };

  const handleKeyDown = (index: number, e: React.KeyboardEvent) => {
    if (e.key === 'Backspace' && !code[index] && index > 0) {
      inputRefs.current[index - 1]?.focus();
    }
  };

  const handlePaste = (e: React.ClipboardEvent) => {
    e.preventDefault();
    const pastedData = e.clipboardData.getData('text').replace(/\D/g, '').slice(0, 6);
    const newCode = [...code];
    for (let i = 0; i < pastedData.length; i++) {
      newCode[i] = pastedData[i];
    }
    setCode(newCode);
    // Focus the next empty input or the last one
    const nextEmptyIndex = newCode.findIndex(c => !c);
    inputRefs.current[nextEmptyIndex === -1 ? 5 : nextEmptyIndex]?.focus();
  };

  const handleVerify = async () => {
    if (!challenge || !selectedMethod) return;

    const verificationCode = selectedMethod === 'BACKUP'
      ? backupCode.trim()
      : code.join('');

    if (selectedMethod !== 'BACKUP' && verificationCode.length !== 6) {
      setError('Please enter a 6-digit code');
      return;
    }

    if (selectedMethod === 'BACKUP' && !verificationCode) {
      setError('Please enter a backup code');
      return;
    }

    setLoading(true);
    setError('');

    try {
      const response = await verifyMFALogin(
        challenge.user_id,
        verificationCode,
        selectedMethod,
        trustDevice
      );

      // Clear MFA challenge from storage
      sessionStorage.removeItem('mfa_challenge');

      // Store tokens
      setCookie('access_token', response.access, 7);
      setCookie('refresh_token', response.refresh, 30);

      // Get redirect info from user
      const user = response.user;
      const currentHostname = window.location.hostname;

      // Determine target subdomain
      const isPlatformUser = ['superuser', 'platform_manager', 'platform_support'].includes(user.role);
      let targetSubdomain: string | null = null;

      if (isPlatformUser) {
        targetSubdomain = 'platform';
      } else if (user.business_subdomain) {
        targetSubdomain = user.business_subdomain;
      }

      // Check if we need to redirect
      const targetHostname = targetSubdomain ? `${targetSubdomain}.${window.location.hostname.split('.').slice(-2).join('.')}` : null;
      const needsRedirect = targetSubdomain && targetHostname && currentHostname !== targetHostname;

      if (needsRedirect && targetSubdomain) {
        const targetUrl = buildSubdomainUrl(targetSubdomain, `/?access_token=${response.access}&refresh_token=${response.refresh}`);
        window.location.href = targetUrl;
        return;
      }

      // Navigate to dashboard
      navigate('/');
    } catch (err: any) {
      setError(err.response?.data?.error || 'Invalid verification code');
    } finally {
      setLoading(false);
    }
  };

  const getMethodIcon = (method: string) => {
    switch (method) {
      case 'SMS':
        return <Smartphone className="h-5 w-5" />;
      case 'TOTP':
        return <Shield className="h-5 w-5" />;
      case 'BACKUP':
        return <Key className="h-5 w-5" />;
      default:
        return null;
    }
  };

  const getMethodLabel = (method: string) => {
    switch (method) {
      case 'SMS':
        return challenge?.phone_last_4
          ? `SMS to ***-***-${challenge.phone_last_4}`
          : 'SMS Code';
      case 'TOTP':
        return 'Authenticator App';
      case 'BACKUP':
        return 'Backup Code';
      default:
        return method;
    }
  };

  if (!challenge) {
    return (
      <div className="min-h-screen flex items-center justify-center bg-gray-50 dark:bg-gray-900">
        <Loader2 className="h-8 w-8 animate-spin text-brand-600" />
      </div>
    );
  }

  return (
    <div className="min-h-screen flex flex-col items-center justify-center bg-gray-50 dark:bg-gray-900 px-4">
      <div className="w-full max-w-md">
        {/* Header */}
        <div className="text-center mb-8">
          <div className="flex justify-center mb-4">
            <div className="p-3 bg-brand-100 dark:bg-brand-900/30 rounded-full">
              <Shield className="h-8 w-8 text-brand-600 dark:text-brand-400" />
            </div>
          </div>
          <h1 className="text-2xl font-bold text-gray-900 dark:text-white">
            Two-Factor Authentication
          </h1>
          <p className="mt-2 text-sm text-gray-600 dark:text-gray-400">
            Enter a verification code to complete login
          </p>
        </div>

        {/* Error Message */}
        {error && (
          <div className="mb-6 rounded-lg bg-red-50 dark:bg-red-900/20 p-4 border border-red-100 dark:border-red-800/50">
            <div className="flex items-center gap-2">
              <AlertCircle className="h-5 w-5 text-red-500 dark:text-red-400 flex-shrink-0" />
              <p className="text-sm text-red-700 dark:text-red-300">{error}</p>
            </div>
          </div>
        )}

        {/* Method Selection */}
        <div className="bg-white dark:bg-gray-800 rounded-xl shadow-sm border border-gray-200 dark:border-gray-700 p-6">
          {/* Method Tabs */}
          {challenge.mfa_methods.length > 1 && (
            <div className="flex gap-2 mb-6">
              {challenge.mfa_methods.map((method) => (
                <button
                  key={method}
                  onClick={() => {
                    setSelectedMethod(method);
                    setCode(['', '', '', '', '', '']);
                    setBackupCode('');
                    setError('');
                  }}
                  className={`flex-1 flex items-center justify-center gap-2 px-4 py-2 rounded-lg text-sm font-medium transition-colors ${
                    selectedMethod === method
                      ? 'bg-brand-100 dark:bg-brand-900/30 text-brand-700 dark:text-brand-300'
                      : 'bg-gray-100 dark:bg-gray-700 text-gray-600 dark:text-gray-400 hover:bg-gray-200 dark:hover:bg-gray-600'
                  }`}
                >
                  {getMethodIcon(method)}
                  <span className="hidden sm:inline">{method === 'TOTP' ? 'App' : method === 'BACKUP' ? 'Backup' : 'SMS'}</span>
                </button>
              ))}
            </div>
          )}

          {/* SMS Method */}
          {selectedMethod === 'SMS' && (
            <div className="space-y-4">
              {!smsSent ? (
                <>
                  <p className="text-sm text-gray-600 dark:text-gray-400 text-center">
                    We'll send a verification code to your phone ending in{' '}
                    <span className="font-medium">{challenge.phone_last_4}</span>
                  </p>
                  <button
                    onClick={handleSendSMS}
                    disabled={loading}
                    className="w-full flex items-center justify-center gap-2 px-4 py-3 bg-brand-600 hover:bg-brand-700 text-white rounded-lg font-medium transition-colors disabled:opacity-50"
                  >
                    {loading ? (
                      <>
                        <Loader2 className="h-5 w-5 animate-spin" />
                        Sending...
                      </>
                    ) : (
                      <>
                        <Smartphone className="h-5 w-5" />
                        Send Code
                      </>
                    )}
                  </button>
                </>
              ) : (
                <>
                  <div className="flex items-center justify-center gap-2 text-green-600 dark:text-green-400 mb-4">
                    <CheckCircle className="h-5 w-5" />
                    <span className="text-sm">Code sent!</span>
                  </div>
                  <div className="flex justify-center gap-2" onPaste={handlePaste}>
                    {code.map((digit, index) => (
                      <input
                        key={index}
                        ref={(el) => { inputRefs.current[index] = el; }}
                        type="text"
                        inputMode="numeric"
                        maxLength={1}
                        value={digit}
                        onChange={(e) => handleCodeChange(index, e.target.value)}
                        onKeyDown={(e) => handleKeyDown(index, e)}
                        className="w-12 h-14 text-center text-2xl font-semibold border-2 border-gray-300 dark:border-gray-600 rounded-lg focus:border-brand-500 focus:ring-brand-500 bg-white dark:bg-gray-700 text-gray-900 dark:text-white"
                        autoFocus={index === 0}
                      />
                    ))}
                  </div>
                  <button
                    onClick={handleSendSMS}
                    disabled={loading}
                    className="text-sm text-brand-600 dark:text-brand-400 hover:underline block mx-auto mt-2"
                  >
                    Resend code
                  </button>
                </>
              )}
            </div>
          )}

          {/* TOTP Method */}
          {selectedMethod === 'TOTP' && (
            <div className="space-y-4">
              <p className="text-sm text-gray-600 dark:text-gray-400 text-center">
                Enter the 6-digit code from your authenticator app
              </p>
              <div className="flex justify-center gap-2" onPaste={handlePaste}>
                {code.map((digit, index) => (
                  <input
                    key={index}
                    ref={(el) => { inputRefs.current[index] = el; }}
                    type="text"
                    inputMode="numeric"
                    maxLength={1}
                    value={digit}
                    onChange={(e) => handleCodeChange(index, e.target.value)}
                    onKeyDown={(e) => handleKeyDown(index, e)}
                    className="w-12 h-14 text-center text-2xl font-semibold border-2 border-gray-300 dark:border-gray-600 rounded-lg focus:border-brand-500 focus:ring-brand-500 bg-white dark:bg-gray-700 text-gray-900 dark:text-white"
                    autoFocus={index === 0}
                  />
                ))}
              </div>
            </div>
          )}

          {/* Backup Code Method */}
          {selectedMethod === 'BACKUP' && (
            <div className="space-y-4">
              <p className="text-sm text-gray-600 dark:text-gray-400 text-center">
                Enter one of your backup codes
              </p>
              <input
                type="text"
                value={backupCode}
                onChange={(e) => setBackupCode(e.target.value.toUpperCase())}
                placeholder="XXXX-XXXX"
                className="w-full text-center text-lg font-mono tracking-wider border-2 border-gray-300 dark:border-gray-600 rounded-lg py-3 focus:border-brand-500 focus:ring-brand-500 bg-white dark:bg-gray-700 text-gray-900 dark:text-white placeholder-gray-400"
                autoFocus
              />
              <p className="text-xs text-gray-500 dark:text-gray-400 text-center">
                Each backup code can only be used once
              </p>
            </div>
          )}

          {/* Trust Device Checkbox */}
          <div className="mt-6 flex items-center gap-2">
            <input
              type="checkbox"
              id="trust-device"
              checked={trustDevice}
              onChange={(e) => setTrustDevice(e.target.checked)}
              className="h-4 w-4 rounded border-gray-300 text-brand-600 focus:ring-brand-500"
            />
            <label htmlFor="trust-device" className="text-sm text-gray-600 dark:text-gray-400">
              Trust this device for 30 days
            </label>
          </div>

          {/* Verify Button */}
          {((selectedMethod === 'SMS' && smsSent) || selectedMethod === 'TOTP' || selectedMethod === 'BACKUP') && (
            <button
              onClick={handleVerify}
              disabled={loading}
              className="w-full mt-6 flex items-center justify-center gap-2 px-4 py-3 bg-brand-600 hover:bg-brand-700 text-white rounded-lg font-medium transition-colors disabled:opacity-50"
            >
              {loading ? (
                <>
                  <Loader2 className="h-5 w-5 animate-spin" />
                  Verifying...
                </>
              ) : (
                'Verify'
              )}
            </button>
          )}
        </div>

        {/* Back to Login */}
        <button
          onClick={() => {
            sessionStorage.removeItem('mfa_challenge');
            navigate('/login');
          }}
          className="mt-6 flex items-center justify-center gap-2 text-sm text-gray-600 dark:text-gray-400 hover:text-gray-900 dark:hover:text-white mx-auto"
        >
          <ArrowLeft className="h-4 w-4" />
          Back to login
        </button>

        {/* Footer */}
        <div className="mt-8 flex justify-center">
          <SmoothScheduleLogo className="h-6 w-6 text-gray-400" />
        </div>
      </div>
    </div>
  );
};

export default MFAVerifyPage;