import { useState, FormEvent } from 'react';
import { Navigate } from 'react-router-dom';
import { toast } from 'sonner';
import { AlertTriangle, Copy, Lock, Loader2, ShieldCheck, UserPlus } from 'lucide-react';
import { useAuth } from '@/hooks/useAuth';
import api, { 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';
import type { TotpSetupResponse } from '@/types';

/** 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;
}

type ScreenMode =
  | 'login'
  | 'setup'          // first-run admin account creation
  | 'register'       // open registration
  | 'totp'           // TOTP challenge after login
  | 'mfa_enforce'    // forced MFA setup after login/register
  | 'force_pw';      // admin-forced password change

type MfaEnforceStep = 'qr' | 'verify' | 'backup_codes';

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

  // ── Shared credential fields ──
  const [username, setUsername] = useState('');
  const [password, setPassword] = useState('');
  const [confirmPassword, setConfirmPassword] = useState('');

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

  // ── Registration mode ──
  const [mode, setMode] = useState<ScreenMode>('login');

  // ── Inline error (423 lockout, 403 disabled, 401 bad creds) ──
  const [loginError, setLoginError] = useState<string | null>(null);

  // ── MFA enforcement setup flow ──
  const [mfaEnforceStep, setMfaEnforceStep] = useState<MfaEnforceStep>('qr');
  const [mfaEnforceQr, setMfaEnforceQr] = useState('');
  const [mfaEnforceSecret, setMfaEnforceSecret] = useState('');
  const [mfaEnforceBackupCodes, setMfaEnforceBackupCodes] = useState<string[]>([]);
  const [mfaEnforceCode, setMfaEnforceCode] = useState('');
  const [isMfaEnforceSetupPending, setIsMfaEnforceSetupPending] = useState(false);
  const [isMfaEnforceConfirmPending, setIsMfaEnforceConfirmPending] = useState(false);

  // ── Forced password change ──
  const [forcedNewPassword, setForcedNewPassword] = useState('');
  const [forcedConfirmPassword, setForcedConfirmPassword] = useState('');
  const [isForcePwPending, setIsForcePwPending] = useState(false);

  // Redirect authenticated users (no pending MFA flows)
  if (!isLoading && authStatus?.authenticated && !mfaSetupRequired && mode !== 'force_pw') {
    return <Navigate to="/dashboard" replace />;
  }

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

  // Derive active screen — hook-driven states override local mode
  const activeMode: ScreenMode = mfaRequired
    ? 'totp'
    : mfaSetupRequired
    ? 'mfa_enforce'
    : isSetup
    ? 'setup'
    : mode;

  // ── Handlers ──

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

    if (isSetup) {
      const err = validatePassword(password);
      if (err) { toast.error(err); return; }
      if (password !== confirmPassword) { toast.error('Passwords do not match'); return; }
      try {
        await setup({ username, password });
      } catch (error) {
        toast.error(getErrorMessage(error, 'Failed to create account'));
      }
      return;
    }

    try {
      const result = await login({ username, password });
      // must_change_password: backend issued session but UI must gate the app
      if ('must_change_password' in result && result.must_change_password) {
        setMode('force_pw');
      }
      // mfaSetupRequired / mfaRequired handled by hook state → activeMode switches automatically
    } catch (error: any) {
      const status = error?.response?.status;
      if (status === 423) {
        setLoginError(error.response.data?.detail || 'Account locked. Try again later.');
      } else if (status === 403) {
        setLoginError(error.response.data?.detail || 'Account is disabled. Contact an administrator.');
      } else {
        setLoginError(getErrorMessage(error, 'Invalid username or password'));
      }
    }
  };

  const handleRegisterSubmit = async (e: FormEvent) => {
    e.preventDefault();
    const err = validatePassword(password);
    if (err) { toast.error(err); return; }
    if (password !== confirmPassword) { toast.error('Passwords do not match'); return; }
    try {
      await register({ username, password });
      // On success useAuth invalidates query → Navigate handles redirect
      // If mfa_setup_required the hook sets mfaSetupRequired → activeMode switches
    } catch (error) {
      toast.error(getErrorMessage(error, 'Registration failed'));
    }
  };

  const handleTotpSubmit = async (e: FormEvent) => {
    e.preventDefault();
    try {
      await verifyTotp({ code: totpCode, isBackup: useBackupCode });
    } catch (error) {
      toast.error(getErrorMessage(error, 'Invalid verification code'));
      setTotpCode('');
    }
  };

  const handleMfaEnforceStart = async () => {
    if (!mfaToken) return;
    setIsMfaEnforceSetupPending(true);
    try {
      const { data } = await api.post<TotpSetupResponse>('/auth/totp/enforce-setup', {
        mfa_token: mfaToken,
      });
      setMfaEnforceQr(data.qr_code_base64);
      setMfaEnforceSecret(data.secret);
      setMfaEnforceBackupCodes(data.backup_codes);
      setMfaEnforceStep('qr');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to begin MFA setup'));
    } finally {
      setIsMfaEnforceSetupPending(false);
    }
  };

  const handleMfaEnforceConfirm = async () => {
    if (!mfaToken || !mfaEnforceCode || mfaEnforceCode.length !== 6) {
      toast.error('Enter a 6-digit code from your authenticator app');
      return;
    }
    setIsMfaEnforceConfirmPending(true);
    try {
      await api.post('/auth/totp/enforce-confirm', {
        mfa_token: mfaToken,
        code: mfaEnforceCode,
      });
      // Backend issued session — show backup codes then redirect
      setMfaEnforceStep('backup_codes');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Invalid code — try again'));
      setMfaEnforceCode('');
    } finally {
      setIsMfaEnforceConfirmPending(false);
    }
  };

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

  const handleForcePwSubmit = async (e: FormEvent) => {
    e.preventDefault();
    const err = validatePassword(forcedNewPassword);
    if (err) { toast.error(err); return; }
    if (forcedNewPassword !== forcedConfirmPassword) {
      toast.error('Passwords do not match');
      return;
    }
    setIsForcePwPending(true);
    try {
      await api.post('/auth/change-password', {
        old_password: password, // retained from original login submission
        new_password: forcedNewPassword,
      });
      toast.success('Password updated — welcome to UMBRA');
      // Auth query still has authenticated:true → Navigate will fire after re-render
      setMode('login');
    } catch (error) {
      toast.error(getErrorMessage(error, 'Failed to change password'));
    } finally {
      setIsForcePwPending(false);
    }
  };

  // ── Render helpers ──

  const renderTotpChallenge = () => (
    <>
      <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(/[^A-Za-z0-9-]/g, '').toUpperCase()
                    : 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>
    </>
  );

  const renderMfaEnforce = () => {
    // Show a loading/start state if QR hasn't been fetched yet
    if (!mfaEnforceQr && mfaEnforceStep !== 'backup_codes') {
      return (
        <>
          <CardHeader>
            <div className="flex items-center gap-3">
              <div className="p-1.5 rounded-md bg-amber-500/10">
                <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
              </div>
              <div>
                <CardTitle>Set Up Two-Factor Authentication</CardTitle>
                <CardDescription>Your account requires MFA before you can continue</CardDescription>
              </div>
            </div>
          </CardHeader>
          <CardContent>
            <p className="text-sm text-muted-foreground mb-4">
              An administrator has required that your account be protected with an authenticator app.
              You'll need an app like Google Authenticator, Authy, or 1Password to continue.
            </p>
            <Button className="w-full" onClick={handleMfaEnforceStart} disabled={isMfaEnforceSetupPending}>
              {isMfaEnforceSetupPending ? (
                <><Loader2 className="h-4 w-4 animate-spin" />Generating QR Code</>
              ) : (
                'Begin Setup'
              )}
            </Button>
          </CardContent>
        </>
      );
    }

    if (mfaEnforceStep === 'backup_codes') {
      return (
        <>
          <CardHeader>
            <div className="flex items-center gap-3">
              <div className="p-1.5 rounded-md bg-amber-500/10">
                <AlertTriangle className="h-4 w-4 text-amber-400" aria-hidden="true" />
              </div>
              <div>
                <CardTitle>Save Your Backup Codes</CardTitle>
                <CardDescription>Store these somewhere safe — they won't be shown again</CardDescription>
              </div>
            </div>
          </CardHeader>
          <CardContent className="space-y-4">
            <p className="text-xs text-muted-foreground">
              These {mfaEnforceBackupCodes.length} codes can each be used once if you lose access to
              your authenticator app. MFA is now active on your account.
            </p>
            <div className="grid grid-cols-2 gap-2 bg-secondary rounded-md p-3">
              {mfaEnforceBackupCodes.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={() => {
                // Session is already issued — redirect to app
                window.location.href = '/dashboard';
              }}
            >
              I've saved my backup codes — Enter UMBRA
            </Button>
          </CardContent>
        </>
      );
    }

    if (mfaEnforceStep === 'qr') {
      return (
        <>
          <CardHeader>
            <div className="flex items-center gap-3">
              <div className="p-1.5 rounded-md bg-amber-500/10">
                <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
              </div>
              <div>
                <CardTitle>Scan QR Code</CardTitle>
                <CardDescription>Add UMBRA to your authenticator app</CardDescription>
              </div>
            </div>
          </CardHeader>
          <CardContent className="space-y-4">
            <div className="flex justify-center">
              <img
                src={`data:image/png;base64,${mfaEnforceQr}`}
                alt="TOTP QR code — scan with your authenticator app"
                className="h-44 w-44 rounded-md border border-border"
              />
            </div>
            <p className="text-xs text-muted-foreground text-center">
              Can't scan? Enter this code manually in your app:
            </p>
            <code className="block text-center text-xs font-mono bg-secondary px-3 py-2 rounded-md tracking-widest break-all">
              {mfaEnforceSecret}
            </code>
            <Button className="w-full" onClick={() => setMfaEnforceStep('verify')}>
              Next: Verify Code
            </Button>
          </CardContent>
        </>
      );
    }

    // verify step
    return (
      <>
        <CardHeader>
          <div className="flex items-center gap-3">
            <div className="p-1.5 rounded-md bg-amber-500/10">
              <ShieldCheck className="h-4 w-4 text-amber-400" aria-hidden="true" />
            </div>
            <div>
              <CardTitle>Verify Your Authenticator</CardTitle>
              <CardDescription>Enter the 6-digit code shown in your app</CardDescription>
            </div>
          </div>
        </CardHeader>
        <CardContent className="space-y-4">
          <div className="space-y-2">
            <Label htmlFor="enforce-code">Verification Code</Label>
            <Input
              id="enforce-code"
              type="text"
              inputMode="numeric"
              maxLength={6}
              placeholder="000000"
              value={mfaEnforceCode}
              onChange={(e) => setMfaEnforceCode(e.target.value.replace(/\D/g, ''))}
              className="text-center tracking-widest text-lg"
              autoFocus
              autoComplete="one-time-code"
              onKeyDown={(e) => { if (e.key === 'Enter') handleMfaEnforceConfirm(); }}
            />
          </div>
          <Button
            className="w-full"
            onClick={handleMfaEnforceConfirm}
            disabled={isMfaEnforceConfirmPending}
          >
            {isMfaEnforceConfirmPending ? (
              <><Loader2 className="h-4 w-4 animate-spin" />Verifying</>
            ) : (
              'Verify & Enable MFA'
            )}
          </Button>
          <button
            type="button"
            onClick={() => setMfaEnforceStep('qr')}
            className="w-full text-center text-xs text-muted-foreground hover:text-foreground transition-colors"
          >
            Back to QR code
          </button>
        </CardContent>
      </>
    );
  };

  const renderLoginOrSetup = () => (
    <>
      <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>
        {loginError && (
          <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'
            )}
          >
            <AlertTriangle className="h-4 w-4 text-red-400 shrink-0" aria-hidden="true" />
            <p className="text-xs text-red-400">{loginError}</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)}
              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)}
              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}
          >
            {isLoginPending || isSetupPending ? (
              <><Loader2 className="h-4 w-4 animate-spin" />Please wait</>
            ) : isSetup ? (
              'Create Account'
            ) : (
              'Sign in'
            )}
          </Button>
        </form>

        {/* Open registration link — only shown on login screen when enabled */}
        {!isSetup && registrationOpen && (
          <div className="mt-4 text-center">
            <button
              type="button"
              onClick={() => {
                setMode('register');
                setUsername('');
                setPassword('');
                setConfirmPassword('');
                setLoginError(null);
              }}
              className="text-xs text-muted-foreground hover:text-foreground transition-colors"
            >
              Don't have an account?{' '}
              <span className="text-accent hover:underline">Create one</span>
            </button>
          </div>
        )}
      </CardContent>
    </>
  );

  const renderRegister = () => (
    <>
      <CardHeader>
        <div className="flex items-center gap-3">
          <div className="p-1.5 rounded-md bg-accent/10">
            <UserPlus className="h-4 w-4 text-accent" aria-hidden="true" />
          </div>
          <div>
            <CardTitle>Create Account</CardTitle>
            <CardDescription>Register for access to UMBRA</CardDescription>
          </div>
        </div>
      </CardHeader>
      <CardContent>
        <form onSubmit={handleRegisterSubmit} className="space-y-4">
          <div className="space-y-2">
            <Label htmlFor="reg-username" required>Username</Label>
            <Input
              id="reg-username"
              type="text"
              value={username}
              onChange={(e) => setUsername(e.target.value)}
              placeholder="Choose a username"
              required
              autoFocus
              autoComplete="username"
            />
          </div>
          <div className="space-y-2">
            <Label htmlFor="reg-password" required>Password</Label>
            <Input
              id="reg-password"
              type="password"
              value={password}
              onChange={(e) => setPassword(e.target.value)}
              placeholder="Create a password"
              required
              autoComplete="new-password"
            />
          </div>
          <div className="space-y-2">
            <Label htmlFor="reg-confirm-password" required>Confirm Password</Label>
            <Input
              id="reg-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={isRegisterPending}>
            {isRegisterPending ? (
              <><Loader2 className="h-4 w-4 animate-spin" />Creating account</>
            ) : (
              'Create Account'
            )}
          </Button>
        </form>
        <div className="mt-4 text-center">
          <button
            type="button"
            onClick={() => {
              setMode('login');
              setUsername('');
              setPassword('');
              setConfirmPassword('');
              setLoginError(null);
            }}
            className="text-xs text-muted-foreground hover:text-foreground transition-colors"
          >
            Already have an account?{' '}
            <span className="text-accent hover:underline">Sign in</span>
          </button>
        </div>
      </CardContent>
    </>
  );

  const renderForcedPasswordChange = () => (
    <>
      <CardHeader>
        <div className="flex items-center gap-3">
          <div className="p-1.5 rounded-md bg-amber-500/10">
            <Lock className="h-4 w-4 text-amber-400" aria-hidden="true" />
          </div>
          <div>
            <CardTitle>Password Change Required</CardTitle>
            <CardDescription>An administrator has reset your password. Please set a new one.</CardDescription>
          </div>
        </div>
      </CardHeader>
      <CardContent>
        <form onSubmit={handleForcePwSubmit} className="space-y-4">
          <div className="space-y-2">
            <Label htmlFor="force-new-pw" required>New Password</Label>
            <Input
              id="force-new-pw"
              type="password"
              value={forcedNewPassword}
              onChange={(e) => setForcedNewPassword(e.target.value)}
              placeholder="Create a new password"
              required
              autoFocus
              autoComplete="new-password"
            />
          </div>
          <div className="space-y-2">
            <Label htmlFor="force-confirm-pw" required>Confirm New Password</Label>
            <Input
              id="force-confirm-pw"
              type="password"
              value={forcedConfirmPassword}
              onChange={(e) => setForcedConfirmPassword(e.target.value)}
              placeholder="Confirm your new 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={isForcePwPending}>
            {isForcePwPending ? (
              <><Loader2 className="h-4 w-4 animate-spin" />Saving</>
            ) : (
              'Set New Password'
            )}
          </Button>
        </form>
      </CardContent>
    </>
  );

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

      {/* Wordmark */}
      <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">
        {activeMode === 'totp' && renderTotpChallenge()}
        {activeMode === 'mfa_enforce' && renderMfaEnforce()}
        {activeMode === 'force_pw' && renderForcedPasswordChange()}
        {activeMode === 'register' && renderRegister()}
        {(activeMode === 'login' || activeMode === 'setup') && renderLoginOrSetup()}
      </Card>
    </div>
  );
}