UMBRA/frontend/src/components/settings/TotpSetupSection.tsx

import { useState, useEffect } from 'react';
import { toast } from 'sonner';
import {
  AlertTriangle,
  Check,
  Copy,
  Loader2,
  ShieldCheck,
} from 'lucide-react';
import { Button } from '@/components/ui/button';
import { Input } from '@/components/ui/input';
import { Label } from '@/components/ui/label';
import { Separator } from '@/components/ui/separator';
import {
  Dialog,
  DialogContent,
  DialogHeader,
  DialogTitle,
  DialogDescription,
  DialogFooter,
  DialogClose,
} from '@/components/ui/dialog';
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card';
import api from '@/lib/api';
import { getErrorMessage } from '@/lib/api';
import type { TotpSetupResponse } from '@/types';

type TotpSetupState = 'idle' | 'setup' | 'confirm' | 'backup_codes' | 'enabled';

export default function TotpSetupSection() {
  // ── Password change state ──
  const [passwordForm, setPasswordForm] = useState({
    oldPassword: '',
    newPassword: '',
    confirmPassword: '',
  });
  const [isChangingPassword, setIsChangingPassword] = useState(false);

  // ── TOTP state ──
  const [totpSetupState, setTotpSetupState] = useState<TotpSetupState>('idle');
  const [qrCodeBase64, setQrCodeBase64] = useState('');
  const [totpSecret, setTotpSecret] = useState('');
  const [totpConfirmCode, setTotpConfirmCode] = useState('');
  const [backupCodes, setBackupCodes] = useState<string[]>([]);
  const [isTotpSetupPending, setIsTotpSetupPending] = useState(false);
  const [isTotpConfirmPending, setIsTotpConfirmPending] = useState(false);

  // ── Disable / Regenerate dialog state ──
  const [disableDialogOpen, setDisableDialogOpen] = useState(false);
  const [regenDialogOpen, setRegenDialogOpen] = useState(false);
  const [dialogPassword, setDialogPassword] = useState('');
  const [dialogCode, setDialogCode] = useState('');
  const [isDialogPending, setIsDialogPending] = useState(false);

  // On mount: check TOTP status to set initial state
  useEffect(() => {
    api
      .get<{ enabled: boolean }>('/auth/totp/status')
      .then(({ data }) => {
        setTotpSetupState(data.enabled ? 'enabled' : 'idle');
      })
      .catch(() => {
        // If endpoint not yet available, default to idle
        setTotpSetupState('idle');
      });
  }, []);

  // ── Password change ──
  const handlePasswordChange = async () => {
    const { oldPassword, newPassword, confirmPassword } = passwordForm;
    if (!oldPassword || !newPassword || !confirmPassword) {
      toast.error('All password fields are required');
      return;
    }
    if (newPassword !== confirmPassword) {
      toast.error('New passwords do not match');
      return;
    }
    if (newPassword.length < 12) {
      toast.error('Password must be at least 12 characters');
      return;
    }
    const hasLetter = /[a-zA-Z]/.test(newPassword);
    const hasNonLetter = /[^a-zA-Z]/.test(newPassword);
    if (!hasLetter || !hasNonLetter) {
      toast.error('Password must contain at least one letter and one non-letter character');
      return;
    }
    setIsChangingPassword(true);
    try {
      await api.post('/auth/change-password', {
        old_password: oldPassword,
        new_password: newPassword,
      });
      setPasswordForm({ oldPassword: '', newPassword: '', confirmPassword: '' });
      toast.success('Password changed successfully');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to change password'));
    } finally {
      setIsChangingPassword(false);
    }
  };

  // ── TOTP setup ──
  const handleBeginTotpSetup = async () => {
    setIsTotpSetupPending(true);
    try {
      const { data } = await api.post<TotpSetupResponse>('/auth/totp/setup');
      setQrCodeBase64(data.qr_code_base64);
      setTotpSecret(data.secret);
      setBackupCodes(data.backup_codes);
      setTotpSetupState('setup');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to begin TOTP setup'));
    } finally {
      setIsTotpSetupPending(false);
    }
  };

  const handleTotpConfirm = async () => {
    if (!totpConfirmCode || totpConfirmCode.length !== 6) {
      toast.error('Enter a 6-digit code');
      return;
    }
    setIsTotpConfirmPending(true);
    try {
      const { data } = await api.post<{ backup_codes: string[] }>('/auth/totp/confirm', {
        code: totpConfirmCode,
      });
      setBackupCodes(data.backup_codes);
      setTotpConfirmCode('');
      setTotpSetupState('backup_codes');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Invalid code — try again'));
    } finally {
      setIsTotpConfirmPending(false);
    }
  };

  const handleCopyBackupCodes = async () => {
    try {
      await navigator.clipboard.writeText(backupCodes.join('\n'));
      toast.success('Backup codes copied');
    } catch {
      toast.error('Failed to copy codes');
    }
  };

  const handleBackupCodesConfirmed = () => {
    setBackupCodes([]);
    setQrCodeBase64('');
    setTotpSecret('');
    setTotpSetupState('enabled');
  };

  // ── Disable TOTP ──
  const handleDisableConfirm = async () => {
    if (!dialogPassword || !dialogCode) {
      toast.error('Password and code are required');
      return;
    }
    setIsDialogPending(true);
    try {
      await api.post('/auth/totp/disable', { password: dialogPassword, code: dialogCode });
      setDisableDialogOpen(false);
      setDialogPassword('');
      setDialogCode('');
      setTotpSetupState('idle');
      toast.success('Two-factor authentication disabled');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to disable TOTP'));
    } finally {
      setIsDialogPending(false);
    }
  };

  // ── Regenerate backup codes ──
  const handleRegenConfirm = async () => {
    if (!dialogPassword || !dialogCode) {
      toast.error('Password and code are required');
      return;
    }
    setIsDialogPending(true);
    try {
      const { data } = await api.post<{ backup_codes: string[] }>(
        '/auth/totp/backup-codes/regenerate',
        { password: dialogPassword, code: dialogCode }
      );
      setBackupCodes(data.backup_codes);
      setRegenDialogOpen(false);
      setDialogPassword('');
      setDialogCode('');
      setTotpSetupState('backup_codes');
      toast.success('New backup codes generated');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to regenerate backup codes'));
    } finally {
      setIsDialogPending(false);
    }
  };

  const closeDialog = () => {
    setDisableDialogOpen(false);
    setRegenDialogOpen(false);
    setDialogPassword('');
    setDialogCode('');
  };

  return (
    <>
      <Card>
        <CardHeader>
          <div className="flex items-center gap-3">
            <div className="p-1.5 rounded-md bg-red-500/10">
              <ShieldCheck className="h-4 w-4 text-red-400" aria-hidden="true" />
            </div>
            <div>
              <CardTitle>Authentication</CardTitle>
              <CardDescription>Manage your password and two-factor authentication</CardDescription>
            </div>
          </div>
        </CardHeader>
        <CardContent className="space-y-6">

          {/* Subsection A: Change Password */}
          <div className="space-y-4">
            <p className="text-sm font-medium">Change Password</p>
            <div className="space-y-3">
              <div className="space-y-2">
                <Label htmlFor="old_password">Current Password</Label>
                <Input
                  id="old_password"
                  type="password"
                  value={passwordForm.oldPassword}
                  onChange={(e) => setPasswordForm({ ...passwordForm, oldPassword: e.target.value })}
                  autoComplete="current-password"
                />
              </div>
              <div className="space-y-2">
                <Label htmlFor="new_password">New Password</Label>
                <Input
                  id="new_password"
                  type="password"
                  value={passwordForm.newPassword}
                  onChange={(e) => setPasswordForm({ ...passwordForm, newPassword: e.target.value })}
                  autoComplete="new-password"
                />
              </div>
              <div className="space-y-2">
                <Label htmlFor="confirm_password">Confirm New Password</Label>
                <Input
                  id="confirm_password"
                  type="password"
                  value={passwordForm.confirmPassword}
                  onChange={(e) => setPasswordForm({ ...passwordForm, confirmPassword: e.target.value })}
                  autoComplete="new-password"
                />
              </div>
              <Button
                type="button"
                onClick={handlePasswordChange}
                disabled={isChangingPassword}
                size="sm"
              >
                {isChangingPassword ? (
                  <><Loader2 className="h-4 w-4 animate-spin" />Saving</>
                ) : (
                  'Change Password'
                )}
              </Button>
            </div>
          </div>

          <Separator />

          {/* Subsection B: TOTP MFA Setup */}
          <div className="space-y-4">
            {totpSetupState === 'idle' && (
              <div className="flex items-center justify-between">
                <div>
                  <p className="text-sm font-medium">Two-Factor Authentication</p>
                  <p className="text-xs text-muted-foreground mt-0.5">
                    Add an extra layer of security with an authenticator app
                  </p>
                </div>
                <Button
                  variant="outline"
                  size="sm"
                  onClick={handleBeginTotpSetup}
                  disabled={isTotpSetupPending}
                >
                  {isTotpSetupPending ? <Loader2 className="h-4 w-4 animate-spin" /> : 'Enable MFA'}
                </Button>
              </div>
            )}

            {totpSetupState === 'setup' && (
              <div className="space-y-4">
                <p className="text-sm font-medium">Scan with your authenticator app</p>
                <div className="flex justify-center">
                  <img
                    src={`data:image/png;base64,${qrCodeBase64}`}
                    alt="TOTP QR code — scan with your authenticator app"
                    className="h-40 w-40 rounded-md border border-border"
                  />
                </div>
                <p className="text-xs text-muted-foreground text-center">
                  Can't scan? Enter this code manually:
                </p>
                <code className="block text-center text-xs font-mono bg-secondary px-3 py-2 rounded-md tracking-widest break-all">
                  {totpSecret}
                </code>
                <Button className="w-full" onClick={() => setTotpSetupState('confirm')}>
                  Next: Verify Code
                </Button>
              </div>
            )}

            {totpSetupState === 'confirm' && (
              <div className="space-y-4">
                <p className="text-sm font-medium">Verify your authenticator app</p>
                <p className="text-xs text-muted-foreground">
                  Enter the 6-digit code shown in your app to confirm setup.
                </p>
                <div className="space-y-2">
                  <Label htmlFor="totp-confirm-code">Verification Code</Label>
                  <Input
                    id="totp-confirm-code"
                    type="text"
                    inputMode="numeric"
                    maxLength={6}
                    placeholder="000000"
                    value={totpConfirmCode}
                    onChange={(e) => setTotpConfirmCode(e.target.value.replace(/\D/g, ''))}
                    className="text-center tracking-widest text-lg"
                    autoFocus
                    autoComplete="one-time-code"
                    onKeyDown={(e) => { if (e.key === 'Enter') handleTotpConfirm(); }}
                  />
                </div>
                <Button
                  className="w-full"
                  onClick={handleTotpConfirm}
                  disabled={isTotpConfirmPending}
                >
                  {isTotpConfirmPending ? (
                    <><Loader2 className="h-4 w-4 animate-spin" />Verifying</>
                  ) : (
                    'Verify & Enable'
                  )}
                </Button>
                <button
                  type="button"
                  onClick={() => setTotpSetupState('setup')}
                  className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
                >
                  Back to QR code
                </button>
              </div>
            )}

            {totpSetupState === 'backup_codes' && (
              <div className="space-y-4">
                <div className="flex items-center gap-2">
                  <AlertTriangle className="h-4 w-4 text-amber-400" aria-hidden="true" />
                  <p className="text-sm font-medium text-amber-400">Save these backup codes now</p>
                </div>
                <p className="text-xs text-muted-foreground">
                  These {backupCodes.length} codes can each be used once if you lose access to your
                  authenticator app. Store them somewhere safe — they will not be shown again.
                </p>
                <div className="grid grid-cols-2 gap-2 bg-secondary rounded-md p-3">
                  {backupCodes.map((code, i) => (
                    <code key={i} className="text-xs font-mono text-foreground text-center py-0.5">
                      {code}
                    </code>
                  ))}
                </div>
                <Button
                  variant="outline"
                  size="sm"
                  onClick={handleCopyBackupCodes}
                  className="w-full gap-2"
                >
                  <Copy className="h-4 w-4" />
                  Copy All Codes
                </Button>
                <Button className="w-full" onClick={handleBackupCodesConfirmed}>
                  I've saved my backup codes
                </Button>
              </div>
            )}

            {totpSetupState === 'enabled' && (
              <div className="flex items-start justify-between gap-3">
                <div>
                  <p className="text-sm font-medium flex items-center gap-2 flex-wrap">
                    Two-Factor Authentication
                    <span className="inline-flex items-center gap-1 text-[10px] px-1.5 py-0.5 rounded-full bg-green-500/10 text-green-400 font-medium uppercase tracking-wide">
                      <Check className="h-3 w-3" aria-hidden="true" />
                      Enabled
                    </span>
                  </p>
                  <p className="text-xs text-muted-foreground mt-0.5">
                    Your account is protected with an authenticator app
                  </p>
                </div>
                <div className="flex gap-2 shrink-0">
                  <Button
                    variant="outline"
                    size="sm"
                    onClick={() => setRegenDialogOpen(true)}
                  >
                    New backup codes
                  </Button>
                  <Button
                    variant="destructive"
                    size="sm"
                    onClick={() => setDisableDialogOpen(true)}
                    className="bg-destructive/10 text-destructive hover:bg-destructive/20 border-destructive/30"
                  >
                    Disable
                  </Button>
                </div>
              </div>
            )}
          </div>
        </CardContent>
      </Card>

      {/* Disable TOTP Dialog */}
      <Dialog open={disableDialogOpen} onOpenChange={closeDialog}>
        <DialogContent>
          <DialogClose onClick={closeDialog} />
          <DialogHeader>
            <DialogTitle>Disable Two-Factor Authentication</DialogTitle>
            <DialogDescription>
              Enter your password and a current authenticator code to disable MFA.
            </DialogDescription>
          </DialogHeader>
          <div className="space-y-3">
            <div className="space-y-2">
              <Label htmlFor="disable-password">Password</Label>
              <Input
                id="disable-password"
                type="password"
                value={dialogPassword}
                onChange={(e) => setDialogPassword(e.target.value)}
                autoComplete="current-password"
              />
            </div>
            <div className="space-y-2">
              <Label htmlFor="disable-code">Authenticator Code</Label>
              <Input
                id="disable-code"
                type="text"
                inputMode="numeric"
                maxLength={6}
                placeholder="000000"
                value={dialogCode}
                onChange={(e) => setDialogCode(e.target.value.replace(/\D/g, ''))}
                className="text-center tracking-widest"
                autoComplete="one-time-code"
              />
            </div>
          </div>
          <DialogFooter>
            <Button variant="outline" size="sm" onClick={closeDialog} disabled={isDialogPending}>
              Cancel
            </Button>
            <Button
              variant="destructive"
              size="sm"
              onClick={handleDisableConfirm}
              disabled={isDialogPending}
              className="gap-2"
            >
              {isDialogPending ? <Loader2 className="h-4 w-4 animate-spin" /> : null}
              Disable MFA
            </Button>
          </DialogFooter>
        </DialogContent>
      </Dialog>

      {/* Regenerate Backup Codes Dialog */}
      <Dialog open={regenDialogOpen} onOpenChange={closeDialog}>
        <DialogContent>
          <DialogClose onClick={closeDialog} />
          <DialogHeader>
            <DialogTitle>Generate New Backup Codes</DialogTitle>
            <DialogDescription>
              Your existing backup codes will be invalidated. Enter your password and a current
              authenticator code to continue.
            </DialogDescription>
          </DialogHeader>
          <div className="space-y-3">
            <div className="space-y-2">
              <Label htmlFor="regen-password">Password</Label>
              <Input
                id="regen-password"
                type="password"
                value={dialogPassword}
                onChange={(e) => setDialogPassword(e.target.value)}
                autoComplete="current-password"
              />
            </div>
            <div className="space-y-2">
              <Label htmlFor="regen-code">Authenticator Code</Label>
              <Input
                id="regen-code"
                type="text"
                inputMode="numeric"
                maxLength={6}
                placeholder="000000"
                value={dialogCode}
                onChange={(e) => setDialogCode(e.target.value.replace(/\D/g, ''))}
                className="text-center tracking-widest"
                autoComplete="one-time-code"
              />
            </div>
          </div>
          <DialogFooter>
            <Button variant="outline" size="sm" onClick={closeDialog} disabled={isDialogPending}>
              Cancel
            </Button>
            <Button
              size="sm"
              onClick={handleRegenConfirm}
              disabled={isDialogPending}
              className="gap-2"
            >
              {isDialogPending ? <Loader2 className="h-4 w-4 animate-spin" /> : null}
              Generate New Codes
            </Button>
          </DialogFooter>
        </DialogContent>
      </Dialog>
    </>
  );
}