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

function parseInviteParams(input) {
    const codeMatch = input.match(/[?&]code=([^&]+)/);
    const keyMatch = input.match(/[?&]key=([^&]+)/);
    if (codeMatch && keyMatch) return { code: codeMatch[1], secret: keyMatch[1] };
    return null;
}

const Register = () => {
    const [username, setUsername] = useState('');
    const [password, setPassword] = useState('');
    const [confirmPassword, setConfirmPassword] = useState('');
    const [error, setError] = useState('');
    const [loading, setLoading] = useState(false);
    const [inviteKeys, setInviteKeys] = useState(null);
    const [inviteLinkInput, setInviteLinkInput] = useState('');
    const [activeInviteCode, setActiveInviteCode] = useState(null);

    const navigate = useNavigate();
    const location = useLocation();
    const convex = useConvex();

    const processInvite = async (code, secret) => {
        if (!window.cryptoAPI) {
            setError("Critical Error: Secure Crypto API missing. Run in Electron.");
            return;
        }

        try {
            const result = await convex.query(api.invites.use, { code });
            if (result.error) throw new Error(result.error);

            const blob = JSON.parse(result.encryptedPayload);
            const decrypted = await window.cryptoAPI.decryptData(blob.c, secret, blob.iv, blob.t);
            const keys = JSON.parse(decrypted);

            console.log('Invite keys decrypted successfully:', Object.keys(keys).length);
            setInviteKeys(keys);
            setActiveInviteCode(code);
            setError('');
        } catch (err) {
            console.error('Invite error:', err);
            setError('Invite verification failed: ' + err.message);
        }
    };

    useEffect(() => {
        const params = new URLSearchParams(location.search);
        const code = params.get('code');
        const secret = params.get('key');

        if (code && secret) {
            console.log('Invite detected in URL');
            processInvite(code, secret);
        }
    }, [location]);

    const handleManualInvite = () => {
        const parsed = parseInviteParams(inviteLinkInput);
        if (parsed) {
            processInvite(parsed.code, parsed.secret);
        } else {
            setError("Invalid invite link format.");
        }
    };

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

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

        setLoading(true);

        try {
            console.log('Starting registration for:', username);

            const salt = await window.cryptoAPI.randomBytes(16);
            const mk = await window.cryptoAPI.randomBytes(32);
            const { dek, dak } = await window.cryptoAPI.deriveAuthKeys(password, salt);
            const encryptedMK = JSON.stringify(await window.cryptoAPI.encryptData(mk, dek));
            const hak = await window.cryptoAPI.sha256(dak);
            const keys = await window.cryptoAPI.generateKeys();

            const encryptedPrivateKeys = JSON.stringify({
                rsa: await window.cryptoAPI.encryptData(keys.rsaPriv, mk),
                ed: await window.cryptoAPI.encryptData(keys.edPriv, mk)
            });

            const data = await convex.mutation(api.auth.createUserWithProfile, {
                username,
                salt,
                encryptedMK,
                hak,
                publicKey: keys.rsaPub,
                signingKey: keys.edPub,
                encryptedPrivateKeys,
                inviteCode: activeInviteCode || undefined
            });

            if (data.error) throw new Error(data.error);
            console.log('Registration successful:', data);

            if (inviteKeys && data.userId) {
                console.log('Uploading invite keys...');
                const batchKeys = await Promise.all(
                    Object.entries(inviteKeys).map(async ([channelId, channelKeyHex]) => {
                        const payload = JSON.stringify({ [channelId]: channelKeyHex });
                        const encryptedKeyBundle = await window.cryptoAPI.publicEncrypt(keys.rsaPub, payload);
                        return { channelId, userId: data.userId, encryptedKeyBundle, keyVersion: 1 };
                    }).map(p => p.catch(err => {
                        console.error('Failed to encrypt key for channel:', err);
                        return null;
                    }))
                );

                const validKeys = batchKeys.filter(Boolean);
                if (validKeys.length > 0) {
                    await convex.mutation(api.channelKeys.uploadKeys, { keys: validKeys });
                    console.log('Uploaded invite keys');
                }
            }

            navigate('/');
        } catch (err) {
            console.error('Registration error:', err);
            setError(err.message);
        } finally {
            setLoading(false);
        }
    };

    return (
        <div className="auth-container">
            <div className="auth-box">
                <div className="auth-header">
                    <h2>Create an Account</h2>
                    <p>Join the secure chat! {inviteKeys ? '(Invite Active)' : ''}</p>
                </div>
                {error && <div style={{ color: 'red', marginBottom: 10, textAlign: 'center' }}>{error}</div>}

                {!inviteKeys && (
                    <div style={{ marginBottom: '15px', display: 'flex' }}>
                        <input
                            type="text"
                            placeholder="Paste Invite Link Here..."
                            value={inviteLinkInput}
                            onChange={(e) => setInviteLinkInput(e.target.value)}
                            style={{ flex: 1, marginRight: '8px' }}
                        />
                        <button type="button" onClick={handleManualInvite} className="auth-button" style={{ width: 'auto' }}>
                            Apply
                        </button>
                    </div>
                )}

                {inviteKeys ? (
                    <form onSubmit={handleRegister}>
                        <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>Password</label>
                            <input
                                type="password"
                                value={password}
                                onChange={(e) => setPassword(e.target.value)}
                                required
                                disabled={loading}
                            />
                        </div>
                        <div className="form-group">
                            <label>Confirm 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 ? 'Generating Keys...' : 'Continue'}
                        </button>
                    </form>
                ) : (
                    <div style={{ textAlign: 'center', marginTop: '20px', color: '#b9bbbe' }}>
                        <p>Registration is Invite-Only.</p>
                        <p style={{ fontSize: '0.9em' }}>Please paste a valid invite link above to proceed.</p>

                        {/* Backdoor for First User */}
                        <p style={{ marginTop: '20px', fontSize: '0.8em', cursor: 'pointer', color: '#7289da' }} onClick={() => setInviteKeys({})}>
                            (First User / Verify Setup)
                        </p>
                    </div>
                )}

                <div className="auth-footer">
                    Already have an account? <Link to="/">Log In</Link>
                </div>
            </div>
        </div>
    );
};

export default Register;