UMBRA/frontend/src/components/auth/LockScreen.tsx

import { useState, FormEvent } from 'react';
import { Navigate } from 'react-router-dom';
import { toast } from 'sonner';
import { Lock, Loader2 } from 'lucide-react';
import { useAuth } from '@/hooks/useAuth';
import { getErrorMessage } from '@/lib/api';
import { Card, CardContent, CardDescription, CardHeader, CardTitle } from '@/components/ui/card';
import { Input } from '@/components/ui/input';
import { Button } from '@/components/ui/button';
import { Label } from '@/components/ui/label';
import { cn } from '@/lib/utils';
import AmbientBackground from './AmbientBackground';

/** Validates password against backend rules: 12-128 chars, at least one letter + one non-letter. */
function validatePassword(password: string): string | null {
  if (password.length < 12) return 'Password must be at least 12 characters';
  if (password.length > 128) return 'Password must be at most 128 characters';
  if (!/[a-zA-Z]/.test(password)) return 'Password must contain at least one letter';
  if (!/[^a-zA-Z]/.test(password)) return 'Password must contain at least one non-letter character';
  return null;
}

export default function LockScreen() {
  const { authStatus, isLoading, login, setup, verifyTotp, mfaRequired, isLoginPending, isSetupPending, isTotpPending } = useAuth();

  // Credentials state (shared across login/setup states)
  const [username, setUsername] = useState('');
  const [password, setPassword] = useState('');
  const [confirmPassword, setConfirmPassword] = useState('');

  // TOTP challenge state
  const [totpCode, setTotpCode] = useState('');
  const [useBackupCode, setUseBackupCode] = useState(false);

  // Lockout handling (HTTP 423)
  const [lockoutMessage, setLockoutMessage] = useState<string | null>(null);

  // Redirect authenticated users immediately
  if (!isLoading && authStatus?.authenticated) {
    return <Navigate to="/dashboard" replace />;
  }

  const isSetup = authStatus?.setup_required === true;

  const handleCredentialSubmit = async (e: FormEvent) => {
    e.preventDefault();
    setLockoutMessage(null);

    if (isSetup) {
      // Setup mode: validate password then create account
      const validationError = validatePassword(password);
      if (validationError) {
        toast.error(validationError);
        return;
      }
      if (password !== confirmPassword) {
        toast.error('Passwords do not match');
        return;
      }
      try {
        await setup({ username, password });
        // useAuth invalidates auth query → Navigate above handles redirect
      } catch (error) {
        toast.error(getErrorMessage(error, 'Failed to create account'));
      }
    } else {
      // Login mode
      try {
        await login({ username, password });
        // If mfaRequired becomes true, the TOTP state renders automatically
        // If not required, useAuth invalidates auth query → Navigate above handles redirect
      } catch (error: any) {
        if (error?.response?.status === 423) {
          const msg = error.response.data?.detail || 'Account locked. Try again later.';
          setLockoutMessage(msg);
        } else {
          toast.error(getErrorMessage(error, 'Invalid username or password'));
        }
      }
    }
  };

  const handleTotpSubmit = async (e: FormEvent) => {
    e.preventDefault();
    try {
      await verifyTotp(totpCode);
      // useAuth invalidates auth query → Navigate above handles redirect
    } catch (error) {
      toast.error(getErrorMessage(error, 'Invalid verification code'));
      setTotpCode('');
    }
  };

  return (
    <div className="relative flex min-h-screen flex-col items-center justify-center bg-background p-4 overflow-hidden">
      <AmbientBackground />

      {/* Wordmark — in flex flow above card */}
      <span className="font-heading text-5xl sm:text-6xl font-bold tracking-tight text-accent mb-10 relative z-10 animate-slide-up">
        UMBRA
      </span>

      {/* Auth card */}
      <Card className="w-full max-w-sm relative z-10 border-border/80 animate-slide-up">
        {mfaRequired ? (
          // State C: TOTP challenge
          <>
            <CardHeader>
              <div className="flex items-center gap-3">
                <div className="p-1.5 rounded-md bg-accent/10">
                  <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
                </div>
                <div>
                  <CardTitle>Two-Factor Authentication</CardTitle>
                  <CardDescription>
                    {useBackupCode
                      ? 'Enter one of your backup codes'
                      : 'Enter the code from your authenticator app'}
                  </CardDescription>
                </div>
              </div>
            </CardHeader>
            <CardContent>
              <form onSubmit={handleTotpSubmit} className="space-y-4">
                <div className="space-y-2">
                  <Label htmlFor="totp-code">
                    {useBackupCode ? 'Backup Code' : 'Authenticator Code'}
                  </Label>
                  <Input
                    id="totp-code"
                    type="text"
                    inputMode={useBackupCode ? 'text' : 'numeric'}
                    pattern={useBackupCode ? undefined : '[0-9]*'}
                    maxLength={useBackupCode ? 9 : 6}
                    value={totpCode}
                    onChange={(e) =>
                      setTotpCode(
                        useBackupCode
                          ? e.target.value.replace(/[^0-9-]/g, '')
                          : e.target.value.replace(/\D/g, '')
                      )
                    }
                    placeholder={useBackupCode ? 'XXXX-XXXX' : '000000'}
                    autoFocus
                    autoComplete="one-time-code"
                    className="text-center text-lg tracking-widest"
                  />
                </div>
                <Button type="submit" className="w-full" disabled={isTotpPending}>
                  {isTotpPending ? (
                    <>
                      <Loader2 className="h-4 w-4 animate-spin" />
                      Verifying
                    </>
                  ) : (
                    'Verify'
                  )}
                </Button>
                <button
                  type="button"
                  onClick={() => {
                    setUseBackupCode(!useBackupCode);
                    setTotpCode('');
                  }}
                  className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
                >
                  {useBackupCode ? 'Use authenticator app instead' : 'Use a backup code instead'}
                </button>
              </form>
            </CardContent>
          </>
        ) : (
          // State A (setup) or State B (login)
          <>
            <CardHeader>
              <div className="flex items-center gap-3">
                <div className="p-1.5 rounded-md bg-accent/10">
                  <Lock className="h-4 w-4 text-accent" aria-hidden="true" />
                </div>
                <div>
                  <CardTitle>{isSetup ? 'Welcome to UMBRA' : 'Sign in'}</CardTitle>
                  <CardDescription>
                    {isSetup
                      ? 'Create your account to get started'
                      : 'Enter your credentials to continue'}
                  </CardDescription>
                </div>
              </div>
            </CardHeader>
            <CardContent>
              {/* Lockout warning banner */}
              {lockoutMessage && (
                <div
                  role="alert"
                  className={cn(
                    'flex items-center gap-2 rounded-md border border-red-500/30',
                    'bg-red-500/10 px-3 py-2 mb-4'
                  )}
                >
                  <Lock className="h-4 w-4 text-red-400 shrink-0" aria-hidden="true" />
                  <p className="text-xs text-red-400">{lockoutMessage}</p>
                </div>
              )}

              <form onSubmit={handleCredentialSubmit} className="space-y-4">
                <div className="space-y-2">
                  <Label htmlFor="username" required>Username</Label>
                  <Input
                    id="username"
                    type="text"
                    value={username}
                    onChange={(e) => { setUsername(e.target.value); setLockoutMessage(null); }}
                    placeholder="Enter username"
                    required
                    autoFocus
                    autoComplete="username"
                  />
                </div>

                <div className="space-y-2">
                  <Label htmlFor="password" required>Password</Label>
                  <Input
                    id="password"
                    type="password"
                    value={password}
                    onChange={(e) => { setPassword(e.target.value); setLockoutMessage(null); }}
                    placeholder={isSetup ? 'Create a password' : 'Enter password'}
                    required
                    autoComplete={isSetup ? 'new-password' : 'current-password'}
                  />
                </div>

                {isSetup && (
                  <div className="space-y-2">
                    <Label htmlFor="confirm-password" required>Confirm Password</Label>
                    <Input
                      id="confirm-password"
                      type="password"
                      value={confirmPassword}
                      onChange={(e) => setConfirmPassword(e.target.value)}
                      placeholder="Confirm your password"
                      required
                      autoComplete="new-password"
                    />
                    <p className="text-xs text-muted-foreground">
                      Must be 12-128 characters with at least one letter and one non-letter.
                    </p>
                  </div>
                )}

                <Button
                  type="submit"
                  className="w-full"
                  disabled={isLoginPending || isSetupPending || !!lockoutMessage}
                >
                  {isLoginPending || isSetupPending ? (
                    <>
                      <Loader2 className="h-4 w-4 animate-spin" />
                      Please wait
                    </>
                  ) : isSetup ? (
                    'Create Account'
                  ) : (
                    'Sign in'
                  )}
                </Button>
              </form>
            </CardContent>
          </>
        )}
      </Card>
    </div>
  );
}