import React, { useState } from 'react';
import { Link } from 'react-router-dom';
import { useConvex } from 'convex/react';
import { usePlatform } from '../platform';
import { api } from '../../../../convex/_generated/api';

const Recovery = () => {
    const [step, setStep] = useState(1); // 1 = verify key, 2 = set new password
    const [username, setUsername] = useState('');
    const [recoveryKey, setRecoveryKey] = useState('');
    const [newPassword, setNewPassword] = useState('');
    const [confirmPassword, setConfirmPassword] = useState('');
    const [error, setError] = useState('');
    const [loading, setLoading] = useState(false);
    const [success, setSuccess] = useState(false);

    // Stored after step 1 verification
    const [verifiedMK, setVerifiedMK] = useState(null);
    const [recoveredSigningKey, setRecoveredSigningKey] = useState(null);

    const convex = useConvex();
    const { crypto } = usePlatform();

    async function decryptEncryptedField(encryptedJson, keyHex) {
        const obj = JSON.parse(encryptedJson);
        return crypto.decryptData(obj.content, keyHex, obj.iv, obj.tag);
    }

    const handleVerify = async (e) => {
        e.preventDefault();
        setError('');
        setLoading(true);

        try {
            // Clean input: strip whitespace, dashes, non-hex chars
            const cleanKey = recoveryKey.replace(/[\s\-]/g, '').toLowerCase();
            if (!/^[0-9a-f]{64}$/.test(cleanKey)) {
                throw new Error('Invalid Recovery Key format. Must be 64 hex characters.');
            }

            const data = await convex.query(api.auth.getRecoveryData, { username });
            if (data.error) {
                throw new Error(data.error);
            }

            // Try decrypting Ed25519 private key with provided MK
            const encryptedPrivateKeysObj = JSON.parse(data.encryptedPrivateKeys);
            let signingKey;
            try {
                signingKey = await decryptEncryptedField(JSON.stringify(encryptedPrivateKeysObj.ed), cleanKey);
            } catch {
                throw new Error('Invalid Recovery Key. Decryption failed.');
            }

            // Store for step 2
            setVerifiedMK(cleanKey);
            setRecoveredSigningKey(signingKey);
            setStep(2);
        } catch (err) {
            setError(err.message);
        } finally {
            setLoading(false);
        }
    };

    const handleReset = async (e) => {
        e.preventDefault();
        setError('');

        if (newPassword !== confirmPassword) {
            setError('Passwords do not match');
            return;
        }

        if (newPassword.length < 1) {
            setError('Password cannot be empty');
            return;
        }

        setLoading(true);

        try {
            // Generate new salt and derive new keys
            const newSalt = await crypto.randomBytes(16);
            const { dek, dak } = await crypto.deriveAuthKeys(newPassword, newSalt);

            // Re-encrypt master key with new DEK
            const newEncryptedMK = JSON.stringify(await crypto.encryptData(verifiedMK, dek));

            // Hash new DAK
            const newHAK = await crypto.sha256(dak);

            // Sign the reset message
            const timestamp = Date.now();
            const message = `password-reset:${username}:${timestamp}`;
            const signature = await crypto.signMessage(recoveredSigningKey, message);

            // Call the server action
            const result = await convex.action(api.recovery.resetPasswordAction, {
                username,
                newSalt,
                newEncryptedMK,
                newHAK,
                signature,
                timestamp,
            });

            if (result.error) {
                throw new Error(result.error);
            }

            setSuccess(true);
        } catch (err) {
            setError(err.message);
        } finally {
            setLoading(false);
        }
    };

    if (success) {
        return (
            <div className="auth-container">
                <div className="auth-box">
                    <div className="auth-header">
                        <h2>Password Reset Complete</h2>
                        <p>Your password has been changed successfully.</p>
                    </div>
                    <div style={{
                        backgroundColor: 'rgba(59, 165, 92, 0.1)', border: '1px solid rgba(59, 165, 92, 0.4)',
                        borderRadius: '4px', padding: '12px', marginBottom: '16px', textAlign: 'center',
                    }}>
                        <div style={{ color: '#3ba55c', fontSize: '14px' }}>
                            You can now log in with your new password. All your encrypted messages remain accessible.
                        </div>
                    </div>
                    <Link to="/" className="auth-button" style={{
                        display: 'block', textAlign: 'center', textDecoration: 'none',
                    }}>
                        Back to Login
                    </Link>
                </div>
            </div>
        );
    }

    return (
        <div className="auth-container">
            <div className="auth-box">
                <div className="auth-header">
                    <h2>Account Recovery</h2>
                    <p>{step === 1
                        ? 'Enter your username and Recovery Key to verify your identity.'
                        : 'Set a new password for your account.'
                    }</p>
                </div>

                {error && <div style={{ color: '#ed4245', marginBottom: 10, textAlign: 'center', fontSize: '14px' }}>{error}</div>}

                {step === 1 ? (
                    <form onSubmit={handleVerify}>
                        <div className="form-group">
                            <label>Username</label>
                            <input
                                type="text"
                                value={username}
                                onChange={(e) => setUsername(e.target.value)}
                                required
                                disabled={loading}
                            />
                        </div>
                        <div className="form-group">
                            <label>Recovery Key</label>
                            <textarea
                                value={recoveryKey}
                                onChange={(e) => setRecoveryKey(e.target.value)}
                                required
                                disabled={loading}
                                placeholder="Paste your 64-character Recovery Key"
                                rows={3}
                                style={{
                                    fontFamily: 'Consolas, "Courier New", monospace',
                                    fontSize: '14px',
                                    resize: 'none',
                                    width: '100%',
                                    boxSizing: 'border-box',
                                }}
                            />
                        </div>
                        <button type="submit" className="auth-button" disabled={loading}>
                            {loading ? 'Verifying...' : 'Verify Recovery Key'}
                        </button>
                    </form>
                ) : (
                    <form onSubmit={handleReset}>
                        <div style={{
                            backgroundColor: 'rgba(59, 165, 92, 0.1)', border: '1px solid rgba(59, 165, 92, 0.4)',
                            borderRadius: '4px', padding: '10px', marginBottom: '16px', textAlign: 'center',
                        }}>
                            <span style={{ color: '#3ba55c', fontSize: '14px' }}>
                                Recovery Key verified for <strong>{username}</strong>
                            </span>
                        </div>
                        <div className="form-group">
                            <label>New Password</label>
                            <input
                                type="password"
                                value={newPassword}
                                onChange={(e) => setNewPassword(e.target.value)}
                                required
                                disabled={loading}
                            />
                        </div>
                        <div className="form-group">
                            <label>Confirm New Password</label>
                            <input
                                type="password"
                                value={confirmPassword}
                                onChange={(e) => setConfirmPassword(e.target.value)}
                                required
                                disabled={loading}
                            />
                        </div>
                        <button type="submit" className="auth-button" disabled={loading}>
                            {loading ? 'Resetting Password...' : 'Reset Password'}
                        </button>
                    </form>
                )}

                <div className="auth-footer">
                    Remember your password? <Link to="/">Log In</Link>
                </div>
            </div>
        </div>
    );
};

export default Recovery;
