src/lib/crypto-service.ts

import { ref, set, get, Database } from 'firebase/database';
import { getDoc, doc, Firestore } from 'firebase/firestore';
import { JsonWebKey } from 'crypto';
import type { ChatRecipient, Group, User, EncryptedMessage } from '@/lib/types';

const isBrowser = typeof window !== 'undefined';

// Cache for group keys, mapping: groupId -> keyVersion -> CryptoKey
const groupKeyCache = new Map<string, Map<number, CryptoKey>>();


async function setupUserKeys(userId: string, rtdb: Database) {
    if (!isBrowser) return;

    const existingKeys = localStorage.getItem(`cryptoKeys_${userId}`);
    if (existingKeys) {
        return;
    }

    try {
        const keyPair = await window.crypto.subtle.generateKey(
            {
                name: 'RSA-OAEP',
                modulusLength: 2048,
                publicExponent: new Uint8Array([1, 0, 1]),
                hash: 'SHA-256',
            },
            true,
            ['encrypt', 'decrypt']
        );

        const publicKey = await window.crypto.subtle.exportKey('jwk', keyPair.publicKey);
        const privateKey = await window.crypto.subtle.exportKey('jwk', keyPair.privateKey);

        await set(ref(rtdb, `users/${userId}/publicKey`), publicKey);

        const keys = { public: publicKey, private: privateKey };
        localStorage.setItem(`cryptoKeys_${userId}`, JSON.stringify(keys));

    } catch (error) {
        console.error("Key generation failed:", error);
    }
}

async function getPublicKey(recipientId: string, rtdb: Database): Promise<JsonWebKey | null> {
    if (!isBrowser) return null;
    const recipientKeyRef = ref(rtdb, `users/${recipientId}/publicKey`);
    const snapshot = await get(recipientKeyRef);
    return snapshot.exists() ? snapshot.val() : null;
}

async function encryptMessage(message: string, recipientId: string, rtdb: Database): Promise<EncryptedMessage> {
    if (!isBrowser) throw new Error("Encryption can only be done in the browser.");

    // 1. Generate a temporary symmetric key (AES)
    const symmetricKey = await window.crypto.subtle.generateKey(
        { name: 'AES-GCM', length: 256 },
        true,
        ['encrypt', 'decrypt']
    );

    // 2. Encrypt the message with the symmetric key
    const iv = window.crypto.getRandomValues(new Uint8Array(12));
    const encodedMessage = new TextEncoder().encode(message);
    const encryptedData = await window.crypto.subtle.encrypt(
        { name: 'AES-GCM', iv },
        symmetricKey,
        encodedMessage
    );

    // 3. Export and encrypt the symmetric key with the recipient's public key (RSA)
    const symmetricKeyJwk = await window.crypto.subtle.exportKey('jwk', symmetricKey);
    const symmetricKeyString = JSON.stringify(symmetricKeyJwk);

    const recipientPublicKeyJwk = await getPublicKey(recipientId, rtdb);
    if (!recipientPublicKeyJwk) throw new Error("Recipient is offline or has no encryption key.");

    const recipientPublicKey = await window.crypto.subtle.importKey(
        'jwk',
        recipientPublicKeyJwk,
        { name: 'RSA-OAEP', hash: 'SHA-256' },
        false,
        ['encrypt']
    );

    const encryptedSymmetricKey = await window.crypto.subtle.encrypt(
        { name: 'RSA-OAEP' },
        recipientPublicKey,
        new TextEncoder().encode(symmetricKeyString)
    );

    // 4. Combine and return
    const ciphertext = btoa(String.fromCharCode(...new Uint8Array(encryptedData)));
    const encryptedKeyString = btoa(String.fromCharCode(...new Uint8Array(encryptedSymmetricKey)));

    return {
        key: encryptedKeyString,
        ciphertext,
        iv: Array.from(iv),
    };
}


async function decryptMessage(encryptedMessage: EncryptedMessage, userId: string): Promise<string> {
    if (!isBrowser) throw new Error("Decryption can only be done in the browser.");

    const keysString = localStorage.getItem(`cryptoKeys_${userId}`);
    if (!keysString) throw new Error(`Cannot decrypt: No local private key found for user ${userId}.`);
    const keys = JSON.parse(keysString);

    // 1. Decrypt the symmetric key with our private key (RSA)
    const privateKey = await window.crypto.subtle.importKey(
        'jwk',
        keys.private,
        { name: 'RSA-OAEP', hash: 'SHA-256' },
        true,
        ['decrypt']
    );

    const encryptedKeyData = Uint8Array.from(atob(encryptedMessage.key), c => c.charCodeAt(0));
    const decryptedKeyString = await window.crypto.subtle.decrypt({ name: 'RSA-OAEP' }, privateKey, encryptedKeyData);
    const symmetricKeyJwk = JSON.parse(new TextDecoder().decode(decryptedKeyString));
    
    const symmetricKey = await window.crypto.subtle.importKey(
        'jwk',
        symmetricKeyJwk,
        { name: 'AES-GCM' },
        true,
        ['decrypt']
    );
    
    // 2. Decrypt the message with the now-decrypted symmetric key
    const iv = new Uint8Array(encryptedMessage.iv!);
    const encryptedData = Uint8Array.from(atob(encryptedMessage.ciphertext), c => c.charCodeAt(0));
    const decrypted = await window.crypto.subtle.decrypt({ name: 'AES-GCM', iv }, symmetricKey, encryptedData);

    return new TextDecoder().decode(decrypted);
}

// --- Group Encryption Functions ---

async function createEncryptedGroupKey(memberUids: string[], rtdb: Database) {
    if (!isBrowser) throw new Error("Group key generation must be done in the browser.");

    const groupKey = await window.crypto.subtle.generateKey(
        { name: 'AES-GCM', length: 256 },
        true,
        ['encrypt', 'decrypt']
    );

    const exportedKey = await window.crypto.subtle.exportKey('jwk', groupKey);
    const groupKeyString = JSON.stringify(exportedKey);

    const encryptedKeys: { [uid: string]: string } = {};

    for (const uid of memberUids) {
        try {
            // Using the hybrid encryption method to encrypt the group key for each member
            const encryptedKeyPayload = await encryptMessage(groupKeyString, uid, rtdb);
            encryptedKeys[uid] = JSON.stringify(encryptedKeyPayload);
        } catch (error) {
            console.warn(`Could not encrypt group key for user ${uid}, they might be offline or without a public key.`);
        }
    }

    return { encryptedKeys, groupKeyString };
}

async function getGroupKey(groupId: string, keyVersion: number, userId: string, db: Firestore): Promise<CryptoKey | null> {
    if (!isBrowser) return null;

    // 1. Check cache first
    const versionCache = groupKeyCache.get(groupId);
    if (versionCache && versionCache.has(keyVersion)) {
        return versionCache.get(keyVersion)!;
    }

    // 2. Fetch from DB if not in cache
    const groupRef = doc(db, 'groups', groupId);
    const groupSnap = await getDoc(groupRef);
    if (!groupSnap.exists()) {
        console.error(`Group ${groupId} not found.`);
        return null;
    }
    const groupData = groupSnap.data() as Group;
    
    const encryptedKeyPayloadString = groupData.encryptedKeys[userId];
    
    if (!encryptedKeyPayloadString) {
        console.warn(`No encrypted key found for user ${userId} in group ${groupId}.`);
        return null;
    }

    try {
        const encryptedKeyPayload: EncryptedMessage = JSON.parse(encryptedKeyPayloadString);
        const decryptedKeyString = await decryptMessage(encryptedKeyPayload, userId);
        const keyJwk = JSON.parse(decryptedKeyString);
        const groupKey = await window.crypto.subtle.importKey(
            'jwk',
            keyJwk,
            { name: 'AES-GCM' },
            true,
            ['encrypt', 'decrypt']
        );
        
        // Update cache
        if (!groupKeyCache.has(groupId)) {
            groupKeyCache.set(groupId, new Map());
        }
        groupKeyCache.get(groupId)!.set(keyVersion, groupKey);

        return groupKey;
    } catch (error) {
        console.error(`Failed to decrypt/import group key for group ${groupId}, version ${keyVersion}:`, error);
        return null;
    }
}

async function encryptGroupMessage(message: string, groupKey: CryptoKey): Promise<string> {
    if (!isBrowser) throw new Error("Encryption must be done in the browser.");

    const iv = window.crypto.getRandomValues(new Uint8Array(12));
    const data = new TextEncoder().encode(message);
    const encrypted = await window.crypto.subtle.encrypt(
        { name: 'AES-GCM', iv },
        groupKey,
        data
    );

    const combined = new Uint8Array(iv.length + encrypted.byteLength);
    combined.set(iv);
    combined.set(new Uint8Array(encrypted), iv.length);

    return btoa(String.fromCharCode(...combined));
}


async function decryptGroupMessage(encryptedMessage: string, groupKey: CryptoKey): Promise<string> {
    if (!isBrowser) throw new Error("Decryption must be done in the browser.");

    const combined = Uint8Array.from(atob(encryptedMessage), c => c.charCodeAt(0));
    const iv = combined.slice(0, 12);
    const data = combined.slice(12);

    const decrypted = await window.crypto.subtle.decrypt(
        { name: 'AES-GCM', iv },
        groupKey,
        data
    );

    return new TextDecoder().decode(decrypted);
}

// --- MEDIA ENCRYPTION ---

async function encryptFile(
    file: File,
    recipient: ChatRecipient,
    senderUid: string,
    rtdb: Database,
    db: Firestore,
    currentGroup?: Group
): Promise<{ encryptedBlob: Blob; encryptionKey: { [uid: string]: string } | { group: string } }> {
    if (!isBrowser) throw new Error("File encryption must be done in the browser.");

    const fileKey = await window.crypto.subtle.generateKey(
        { name: 'AES-GCM', length: 256 },
        true,
        ['encrypt', 'decrypt']
    );

    const iv = window.crypto.getRandomValues(new Uint8Array(12));
    const fileBuffer = await file.arrayBuffer();
    const encryptedFileBuffer = await window.crypto.subtle.encrypt(
        { name: 'AES-GCM', iv },
        fileKey,
        fileBuffer
    );
    
    const combinedBuffer = new Uint8Array(iv.length + encryptedFileBuffer.byteLength);
    combinedBuffer.set(iv);
    combinedBuffer.set(new Uint8Array(encryptedFileBuffer), iv.length);
    const encryptedBlob = new Blob([combinedBuffer]);

    const fileKeyJwk = await window.crypto.subtle.exportKey('jwk', fileKey);
    const fileKeyString = JSON.stringify(fileKeyJwk);

    if (recipient.isGroup && currentGroup) {
        if (!currentGroup.keyLastRotatedAt) throw new Error("Group key version is missing.");
        const groupKey = await getGroupKey(currentGroup.id, currentGroup.keyLastRotatedAt, senderUid, db);
        if (!groupKey) throw new Error("Cannot encrypt file: Group key not available.");
        
        const encryptedFileKey = await encryptGroupMessage(fileKeyString, groupKey);
        return { encryptedBlob, encryptionKey: { group: encryptedFileKey } };

    } else if (!recipient.isGroup) {
        const encryptionKey: { [uid: string]: string } = {};
        const participants = [senderUid, recipient.uid];
        for (const uid of participants) {
            // NOTE: Encrypting the file key string directly. This remains unchanged.
            const recipientPublicKeyJwk = await getPublicKey(uid, rtdb);
            if (!recipientPublicKeyJwk) throw new Error(`Recipient ${uid} has no public key.`);
            const recipientPublicKey = await window.crypto.subtle.importKey('jwk', recipientPublicKeyJwk, { name: 'RSA-OAEP', hash: 'SHA-256' }, false, ['encrypt']);
            const encryptedKey = await window.crypto.subtle.encrypt({ name: 'RSA-OAEP' }, recipientPublicKey, new TextEncoder().encode(fileKeyString));
            encryptionKey[uid] = btoa(String.fromCharCode(...new Uint8Array(encryptedKey)));
        }
        return { encryptedBlob, encryptionKey };
    }
    
    throw new Error("Invalid recipient type for file encryption.");
}

async function decryptFile(
    encryptedBlob: Blob,
    encryptionKey: { [uid: string]: string } | { group: string },
    keyVersion: number | undefined,
    currentUser: User,
    db: Firestore,
    group?: Group,
    originalFileType?: string
): Promise<Blob> {
    if (!isBrowser) throw new Error("File decryption must be done in the browser.");
    
    try {
        let decryptedKeyString: string;

        // Decrypt the symmetric key for the file.
        if (group && 'group' in encryptionKey) {
            const currentKeyVersion = group.keyLastRotatedAt;
            if (!currentKeyVersion) throw new Error("Group is missing a current key version.");

            const groupKey = await getGroupKey(group.id, currentKeyVersion, currentUser.uid, db);
            if (!groupKey) throw new Error("Decryption failed: Key is outdated or user lacks permission.");

            decryptedKeyString = await decryptGroupMessage(encryptionKey.group, groupKey);
        } else if ('group' in encryptionKey) {
            throw new Error("Media is group encrypted, but no group context provided.");
        } else {
            const encryptedKeyForUser = encryptionKey[currentUser.uid];
            if (!encryptedKeyForUser) throw new Error("Missing personal encryption key for media.");
            
            // NOTE: Using direct RSA decryption for the file key string. This remains unchanged.
            const keysString = localStorage.getItem(`cryptoKeys_${currentUser.uid}`);
            if (!keysString) throw new Error(`No local private key for user ${currentUser.uid}.`);
            const keys = JSON.parse(keysString);
            const privateKey = await window.crypto.subtle.importKey('jwk', keys.private, { name: 'RSA-OAEP', hash: 'SHA-256' }, true, ['decrypt']);
            const encryptedData = Uint8Array.from(atob(encryptedKeyForUser), c => c.charCodeAt(0));
            const decrypted = await window.crypto.subtle.decrypt({ name: 'RSA-OAEP' }, privateKey, encryptedData);
            decryptedKeyString = new TextDecoder().decode(decrypted);
        }

        const fileKeyJwk = JSON.parse(decryptedKeyString);
        const fileKey = await window.crypto.subtle.importKey(
            'jwk',
            fileKeyJwk,
            { name: 'AES-GCM' },
            true,
            ['decrypt']
        );
        
        const encryptedBuffer = await encryptedBlob.arrayBuffer();
        const iv = encryptedBuffer.slice(0, 12);
        const data = encryptedBuffer.slice(12);

        const decryptedFileBuffer = await window.crypto.subtle.decrypt(
            { name: 'AES-GCM', iv: new Uint8Array(iv) },
            fileKey,
            data
        );
        
        // If an original file type is provided, use it to reconstruct the Blob correctly.
        if (originalFileType) {
            return new Blob([decryptedFileBuffer], { type: originalFileType });
        }

        return new Blob([decryptedFileBuffer]);

    } catch (error: any) {
         console.error("File decryption failed:", error);
         if (error.message.includes("Decryption failed: Key is outdated")) {
             throw error;
         }
         throw new Error("Could not decrypt media file.");
    }
}


export const cryptoService = {
    setupUserKeys,
    getPublicKey,
    encryptMessage,
    decryptMessage,
    createEncryptedGroupKey,
    getGroupKey,
    encryptGroupMessage,
    decryptGroupMessage,
    encryptFile,
    decryptFile,
    storeGroupKey: (groupId: string, keyVersion: number, keyString: string) => {
        if (!isBrowser) return;
        try {
            const keyJwk = JSON.parse(keyString);
            window.crypto.subtle.importKey('jwk', keyJwk, { name: 'AES-GCM' }, true, ['encrypt', 'decrypt'])
                .then(key => {
                    if (!groupKeyCache.has(groupId)) {
                        groupKeyCache.set(groupId, new Map());
                    }
                    groupKeyCache.get(groupId)!.set(keyVersion, key);
                });
        } catch (error) {
            console.error("Failed to store group key in cache", error);
        }
    },
    clearGroupKeyCache: (groupId: string) => {
        if (!isBrowser) return;
        groupKeyCache.delete(groupId);
    }
};