import { useState, FormEvent } from 'react';
import { Link, useNavigate, useLocation } from 'react-router-dom';
import { Shield, Lock } from 'lucide-react';
import api from '../lib/api';
import { useAuthStore } from '../lib/store';
import { deriveKeyFromPassword, decryptMasterKey, decryptPrivateKeyToCryptoKey } from '../crypto';
import { socket } from '../lib/ws';

export default function LoginPage() {
  const navigate = useNavigate();
  const location = useLocation();
  const { setAuth, setMasterKey, setKeyPair, setUsername, username, email: storedEmail } = useAuthStore();

  // Reauth: tokens still valid but crypto keys lost on refresh
  const isReauth = location.state?.reauth === true && !!storedEmail;

  const [email, setEmail] = useState('');
  const [password, setPassword] = useState('');
  const [totpCode, setTotpCode] = useState('');
  const [require2fa, setRequire2fa] = useState(false);
  const [loading, setLoading] = useState(false);
  const [error, setError] = useState('');

  async function submit(e: FormEvent) {
    e.preventDefault();
    setError('');
    setLoading(true);

    try {
      const res = await api.post('/auth/login', {
        email: isReauth ? storedEmail : email,
        password,
        ...(totpCode ? { totpCode } : {}),
      });

      if (res.data.require2fa) {
        setRequire2fa(true);
        setLoading(false);
        return;
      }

      const { accessToken, refreshToken, userId, username: uname, kdfSalt, encryptedMasterKey, isAdmin } = res.data;

      const derived = await deriveKeyFromPassword(password, kdfSalt);
      const masterKey = await decryptMasterKey(encryptedMasterKey, derived);

      setAuth({ accessToken, refreshToken, userId, isAdmin, email: isReauth ? storedEmail! : email });
      setMasterKey(masterKey);
      setUsername(uname);

      const keyRes = await api.get('/auth/my-key');
      const privateKey = await decryptPrivateKeyToCryptoKey(keyRes.data.encrypted_private_key, masterKey);
      setKeyPair(keyRes.data.public_key, privateKey);

      socket.connect(accessToken);
      navigate('/messages');
    } catch (err: any) {
      setError(err.response?.data?.error || 'Login failed');
    } finally {
      setLoading(false);
    }
  }

  return (
    <div className="min-h-screen flex items-center justify-center p-4">
      <div className="w-full max-w-sm">
        <div className="flex flex-col items-center mb-8">
          {isReauth ? (
            <Lock className="w-12 h-12 text-accent mb-3" />
          ) : (
            <Shield className="w-12 h-12 text-accent mb-3" />
          )}
          <h1 className="text-2xl font-bold">CloudLink</h1>
          <p className="text-slate-400 text-sm mt-1">
            {isReauth ? `Welcome back, ${username}` : 'End-to-end encrypted'}
          </p>
        </div>

        <form onSubmit={submit} className="bg-surface-1 rounded-xl p-6 space-y-4 border border-slate-700">
          <h2 className="font-semibold text-lg">
            {isReauth ? 'Enter your password to unlock' : 'Sign in'}
          </h2>

          {error && <p className="text-red-400 text-sm">{error}</p>}

          {isReauth ? (
            <Field label="Password" type="password" value={password} onChange={setPassword} autoFocus />
          ) : !require2fa ? (
            <>
              <Field label="Email" type="email" value={email} onChange={setEmail} />
              <Field label="Password" type="password" value={password} onChange={setPassword} />
            </>
          ) : (
            <div>
              <label className="block text-sm text-slate-400 mb-1">Authenticator code</label>
              <input
                className="input w-full"
                type="text"
                inputMode="numeric"
                maxLength={6}
                placeholder="000000"
                value={totpCode}
                onChange={(e) => setTotpCode(e.target.value)}
                autoFocus
              />
            </div>
          )}

          <button type="submit" disabled={loading} className="btn-primary w-full">
            {loading ? (isReauth ? 'Unlocking…' : 'Signing in…') : isReauth ? 'Unlock' : require2fa ? 'Verify' : 'Sign in'}
          </button>

          {isReauth && (
            <button
              type="button"
              onClick={() => { useAuthStore.getState().logout(); navigate('/login'); }}
              className="w-full text-center text-sm text-slate-500 hover:text-slate-300"
            >
              Sign in as a different user
            </button>
          )}
        </form>

        {!isReauth && (
          <p className="text-center text-sm text-slate-400 mt-4">
            No account?{' '}
            <Link to="/register" className="text-accent hover:underline">Register</Link>
          </p>
        )}
      </div>
    </div>
  );
}

function Field({ label, type, value, onChange, autoFocus }: {
  label: string; type: string; value: string; onChange: (v: string) => void; autoFocus?: boolean;
}) {
  return (
    <div>
      <label className="block text-sm text-slate-400 mb-1">{label}</label>
      <input
        className="input w-full"
        type={type}
        value={value}
        onChange={(e) => onChange(e.target.value)}
        autoFocus={autoFocus}
        required
      />
    </div>
  );
}