// SpotNotes End-to-End Encryption (E2EE) Module
// Provides client-side zero-knowledge encryption using standard Web Crypto API.
// Server operator never has access to plaintext private notes or user private keys.

window.SpotCrypto = window.SpotCrypto || {};

(function() {
    const PBKDF2_ITERATIONS = 100000;
    const SESSION_KEY_NAME = 'spotnotes_user_privkey';

    // In-memory cache for decrypted symmetric notepad keys: padId -> CryptoKey
    const padKeyCache = new Map();
    let currentUserPrivateKey = null;

    // --- Utility: Encoding / Decoding ---
    function bufToB64(buffer) {
        let binary = '';
        const bytes = new Uint8Array(buffer);
        for (let i = 0; i < bytes.byteLength; i++) {
            binary += String.fromCharCode(bytes[i]);
        }
        return window.btoa(binary);
    }

    function b64ToBuf(base64) {
        const binary = window.atob(base64);
        const bytes = new Uint8Array(binary.length);
        for (let i = 0; i < binary.length; i++) {
            bytes[i] = binary.charCodeAt(i);
        }
        return bytes.buffer;
    }

    function bufToB64Url(buffer) {
        return bufToB64(buffer).replace(/\+/g, '-').replace(/\//g, '_').replace(/=+$/, '');
    }

    function b64UrlToBuf(b64url) {
        let b64 = b64url.replace(/-/g, '+').replace(/_/g, '/');
        while (b64.length % 4) b64 += '=';
        return b64ToBuf(b64);
    }

    // --- Key Derivation (PBKDF2) for User Password ---
    async function deriveKEK(password, saltUint8Array) {
        const enc = new TextEncoder();
        const baseKey = await crypto.subtle.importKey(
            'raw',
            enc.encode(password),
            'PBKDF2',
            false,
            ['deriveKey']
        );
        return await crypto.subtle.deriveKey(
            {
                name: 'PBKDF2',
                salt: saltUint8Array,
                iterations: PBKDF2_ITERATIONS,
                hash: 'SHA-256'
            },
            baseKey,
            { name: 'AES-GCM', length: 256 },
            false,
            ['encrypt', 'decrypt']
        );
    }

    // --- User Keypair (RSA-OAEP) Generation & Storage ---
    window.SpotCrypto.generateUserRegistrationPayload = async function(email, password, firstName, lastName) {
        // 1. Generate RSA-OAEP Keypair
        const keyPair = await crypto.subtle.generateKey(
            {
                name: 'RSA-OAEP',
                modulusLength: 2048,
                publicExponent: new Uint8Array([1, 0, 1]),
                hash: 'SHA-256'
            },
            true,
            ['encrypt', 'decrypt']
        );

        // 2. Export Public Key as JWK string
        const pubJwk = await crypto.subtle.exportKey('jwk', keyPair.publicKey);
        const publicKeyStr = JSON.stringify(pubJwk);

        // 3. Generate random 16-byte salt for user password KEK
        const salt = crypto.getRandomValues(new Uint8Array(16));
        const saltStr = bufToB64(salt.buffer);

        // 4. Derive KEK and encrypt the private key
        const kek = await deriveKEK(password, salt);
        const privJwk = await crypto.subtle.exportKey('jwk', keyPair.privateKey);
        const privJsonBytes = new TextEncoder().encode(JSON.stringify(privJwk));

        const iv = crypto.getRandomValues(new Uint8Array(12));
        const encryptedPrivBuffer = await crypto.subtle.encrypt(
            { name: 'AES-GCM', iv: iv },
            kek,
            privJsonBytes
        );

        const encPrivatePayload = JSON.stringify({
            iv: bufToB64(iv.buffer),
            ct: bufToB64(encryptedPrivBuffer)
        });

        // Cache private key in session
        currentUserPrivateKey = keyPair.privateKey;
        try {
            sessionStorage.setItem(SESSION_KEY_NAME, JSON.stringify(privJwk));
        } catch (e) { console.warn('sessionStorage unavailable', e); }

        return {
            email,
            password,
            first_name: firstName,
            last_name: lastName,
            public_key: publicKeyStr,
            encrypted_private_key: encPrivatePayload,
            encryption_salt: saltStr
        };
    };

    // Unlock User Private Key on Login
    window.SpotCrypto.unlockUserPrivateKey = async function(user, password) {
        if (!user.encrypted_private_key || !user.encryption_salt) {
            console.warn('User has no encryption key stored.');
            return null;
        }

        try {
            const salt = new Uint8Array(b64ToBuf(user.encryption_salt));
            const kek = await deriveKEK(password, salt);

            const encBundle = JSON.parse(user.encrypted_private_key);
            const iv = new Uint8Array(b64ToBuf(encBundle.iv));
            const ct = b64ToBuf(encBundle.ct);

            const decryptedBuf = await crypto.subtle.decrypt(
                { name: 'AES-GCM', iv: iv },
                kek,
                ct
            );

            const privJwk = JSON.parse(new TextDecoder().decode(decryptedBuf));
            const privateKey = await crypto.subtle.importKey(
                'jwk',
                privJwk,
                { name: 'RSA-OAEP', hash: 'SHA-256' },
                true,
                ['decrypt']
            );

            currentUserPrivateKey = privateKey;
            try {
                sessionStorage.setItem(SESSION_KEY_NAME, JSON.stringify(privJwk));
            } catch (e) { console.warn('sessionStorage unavailable', e); }

            return privateKey;
        } catch (err) {
            console.error('Failed to unlock private key:', err);
            throw new Error('Failed to decrypt user private encryption key.');
        }
    };

    // Restore Private Key from Session Storage (e.g., page refresh)
    window.SpotCrypto.restoreUserPrivateKeyFromSession = async function() {
        if (currentUserPrivateKey) return currentUserPrivateKey;

        try {
            const stored = sessionStorage.getItem(SESSION_KEY_NAME);
            if (!stored) return null;
            const privJwk = JSON.parse(stored);
            currentUserPrivateKey = await crypto.subtle.importKey(
                'jwk',
                privJwk,
                { name: 'RSA-OAEP', hash: 'SHA-256' },
                true,
                ['decrypt']
            );
            return currentUserPrivateKey;
        } catch (e) {
            console.warn('Could not restore private key from session', e);
            return null;
        }
    };

    window.SpotCrypto.clearSession = function() {
        currentUserPrivateKey = null;
        padKeyCache.clear();
        try {
            sessionStorage.removeItem(SESSION_KEY_NAME);
        } catch (e) {}
    };

    // --- Notepad Encryption Keys (AES-GCM 256) ---

    // Generate a fresh random AES-GCM 256 key for a new private notepad
    window.SpotCrypto.createNotepadKey = async function(ownerPublicKeyJwkStr) {
        // 1. Generate AES-GCM key
        const padKey = await crypto.subtle.generateKey(
            { name: 'AES-GCM', length: 256 },
            true,
            ['encrypt', 'decrypt']
        );

        // 2. Export raw key bytes
        const rawKey = await crypto.subtle.exportKey('raw', padKey);

        // 3. Encrypt raw key with Owner's RSA Public Key
        const pubJwk = JSON.parse(ownerPublicKeyJwkStr);
        const ownerPubKey = await crypto.subtle.importKey(
            'jwk',
            pubJwk,
            { name: 'RSA-OAEP', hash: 'SHA-256' },
            false,
            ['encrypt']
        );

        const wrappedKeyBuf = await crypto.subtle.encrypt(
            { name: 'RSA-OAEP' },
            ownerPubKey,
            rawKey
        );

        const encryptedKeyB64 = bufToB64(wrappedKeyBuf);
        return {
            padKey,
            rawKey,
            encryptedKeyB64
        };
    };

    // Encrypt an existing Notepad Key for a collaborator using their RSA public key
    window.SpotCrypto.wrapPadKeyForCollaborator = async function(padKey, collaboratorPublicKeyJwkStr) {
        const rawKey = await crypto.subtle.exportKey('raw', padKey);
        const pubJwk = typeof collaboratorPublicKeyJwkStr === 'string' ? JSON.parse(collaboratorPublicKeyJwkStr) : collaboratorPublicKeyJwkStr;
        const collabPubKey = await crypto.subtle.importKey(
            'jwk',
            pubJwk,
            { name: 'RSA-OAEP', hash: 'SHA-256' },
            false,
            ['encrypt']
        );

        const wrappedKeyBuf = await crypto.subtle.encrypt(
            { name: 'RSA-OAEP' },
            collabPubKey,
            rawKey
        );
        return bufToB64(wrappedKeyBuf);
    };

    // Decrypt a wrapped Notepad Key using user's private key
    window.SpotCrypto.unwrapNotepadKey = async function(padId, encryptedKeyB64) {
        if (!encryptedKeyB64) return null;
        if (padKeyCache.has(padId)) return padKeyCache.get(padId);

        const privKey = await window.SpotCrypto.restoreUserPrivateKeyFromSession();
        if (!privKey) {
            console.warn(`No private key available to decrypt notepad ${padId}`);
            return null;
        }

        try {
            const wrappedBuf = b64ToBuf(encryptedKeyB64);
            const rawKeyBuf = await crypto.subtle.decrypt(
                { name: 'RSA-OAEP' },
                privKey,
                wrappedBuf
            );

            const padKey = await crypto.subtle.importKey(
                'raw',
                rawKeyBuf,
                { name: 'AES-GCM', length: 256 },
                true,
                ['encrypt', 'decrypt']
            );

            padKeyCache.set(padId, padKey);
            return padKey;
        } catch (err) {
            console.error(`Failed to unwrap key for notepad ${padId}:`, err);
            return null;
        }
    };

    // Export a Notepad Key to base64url for zero-knowledge Link Sharing (#key=...)
    window.SpotCrypto.exportPadKeyForLink = async function(padKey) {
        const rawKey = await crypto.subtle.exportKey('raw', padKey);
        return bufToB64Url(rawKey);
    };

    // Import a Notepad Key from a link URL fragment (#key=...)
    window.SpotCrypto.importPadKeyFromLink = async function(padId, keyB64Url) {
        try {
            const rawKeyBuf = b64UrlToBuf(keyB64Url);
            const padKey = await crypto.subtle.importKey(
                'raw',
                rawKeyBuf,
                { name: 'AES-GCM', length: 256 },
                true,
                ['encrypt', 'decrypt']
            );
            padKeyCache.set(padId, padKey);
            return padKey;
        } catch (err) {
            console.error('Failed to import pad key from link hash:', err);
            return null;
        }
    };

    window.SpotCrypto.getCachedPadKey = function(padId) {
        return padKeyCache.get(padId) || null;
    };

    window.SpotCrypto.setCachedPadKey = function(padId, key) {
        padKeyCache.set(padId, key);
    };

    // --- Note Content Encryption & Decryption ---

    // Encrypts { title, content, lat, lng } client-side
    window.SpotCrypto.encryptNote = async function(noteData, padKey) {
        const payload = {
            title: noteData.title || '',
            content: noteData.content || '',
            lat: parseFloat(noteData.lat) || 0.0,
            lng: parseFloat(noteData.lng) || 0.0,
            ts: Date.now()
        };

        const jsonBytes = new TextEncoder().encode(JSON.stringify(payload));
        const iv = crypto.getRandomValues(new Uint8Array(12));

        const ciphertextBuf = await crypto.subtle.encrypt(
            { name: 'AES-GCM', iv: iv },
            padKey,
            jsonBytes
        );

        return `ENC1:${bufToB64(iv.buffer)}:${bufToB64(ciphertextBuf)}`;
    };

    // Decrypts a note client-side. If not encrypted, returns note as-is.
    window.SpotCrypto.decryptNote = async function(note, padKey) {
        if (!note || !note.is_encrypted) {
            return {
                ...note,
                title: note.title || '',
                content: note.content || '',
                lat: note.lat,
                lng: note.lng
            };
        }

        if (!padKey) {
            return {
                ...note,
                title: '🔒 [Encrypted Private Note]',
                content: 'Key required to decrypt this note.',
                is_locked: true
            };
        }

        try {
            const rawContent = note.content || '';
            if (!rawContent.startsWith('ENC1:')) {
                return {
                    ...note,
                    title: '🔒 [Encrypted Private Note]',
                    content: 'Unrecognized encryption format.'
                };
            }

            const parts = rawContent.split(':');
            const iv = new Uint8Array(b64ToBuf(parts[1]));
            const ct = b64ToBuf(parts[2]);

            const decryptedBuf = await crypto.subtle.decrypt(
                { name: 'AES-GCM', iv: iv },
                padKey,
                ct
            );

            const payload = JSON.parse(new TextDecoder().decode(decryptedBuf));
            return {
                ...note,
                title: payload.title,
                content: payload.content,
                lat: payload.lat,
                lng: payload.lng,
                is_locked: false
            };
        } catch (err) {
            console.warn(`Failed to decrypt note ${note.id}:`, err);
            return {
                ...note,
                title: '🔒 [Decryption Error]',
                content: 'Failed to decrypt note content with current key.',
                is_locked: true
            };
        }
    };

    // Decrypt an array of notes
    window.SpotCrypto.decryptNotes = async function(notes, padKey) {
        if (!notes || !Array.isArray(notes)) return [];
        return Promise.all(notes.map(n => window.SpotCrypto.decryptNote(n, padKey)));
    };
})();