/**
 * End-to-End Encryption Utilities
 * Client-side encryption using Web Crypto API
 * - RSA-OAEP for key exchange
 * - AES-256-GCM for message content
 */

/**
 * Generate RSA key pair for user
 * @returns {Promise<{publicKey: string, privateKey: string}>}
 */
export async function generateKeyPair() {
  const keyPair = await window.crypto.subtle.generateKey(
    {
      name: 'RSA-OAEP',
      modulusLength: 2048,
      publicExponent: new Uint8Array([1, 0, 1]),
      hash: 'SHA-256'
    },
    true, // extractable
    ['encrypt', 'decrypt']
  );

  // Export keys
  const publicKey = await window.crypto.subtle.exportKey('spki', keyPair.publicKey);
  const privateKey = await window.crypto.subtle.exportKey('pkcs8', keyPair.privateKey);

  return {
    publicKey: arrayBufferToBase64(publicKey),
    privateKey: arrayBufferToBase64(privateKey)
  };
}

/**
 * Store private key encrypted with user's password
 * @param {string} privateKey - Base64 encoded private key
 * @param {string} password - User's password
 */
export async function storePrivateKey(privateKey, password) {
  // Derive encryption key from password
  const passwordKey = await deriveKeyFromPassword(password);
  
  // Encrypt private key
  const iv = window.crypto.getRandomValues(new Uint8Array(12));
  const encrypted = await window.crypto.subtle.encrypt(
    {
      name: 'AES-GCM',
      iv: iv
    },
    passwordKey,
    base64ToArrayBuffer(privateKey)
  );

  // Store in localStorage
  localStorage.setItem('encryptedPrivateKey', arrayBufferToBase64(encrypted));
  localStorage.setItem('privateKeyIV', arrayBufferToBase64(iv));
}

/**
 * Retrieve and decrypt private key
 * @param {string} password - User's password
 * @returns {Promise<CryptoKey>} Decrypted private key
 */
export async function getPrivateKey(password) {
  const encryptedKey = localStorage.getItem('encryptedPrivateKey');
  const iv = localStorage.getItem('privateKeyIV');

  if (!encryptedKey || !iv) {
    throw new Error('No stored private key found');
  }

  // Derive decryption key from password
  const passwordKey = await deriveKeyFromPassword(password);

  // Decrypt private key
  const decrypted = await window.crypto.subtle.decrypt(
    {
      name: 'AES-GCM',
      iv: base64ToArrayBuffer(iv)
    },
    passwordKey,
    base64ToArrayBuffer(encryptedKey)
  );

  // Import as CryptoKey
  return await window.crypto.subtle.importKey(
    'pkcs8',
    decrypted,
    {
      name: 'RSA-OAEP',
      hash: 'SHA-256'
    },
    true,
    ['decrypt']
  );
}

/**
 * Derive encryption key from password using PBKDF2
 * @param {string} password - User's password
 * @returns {Promise<CryptoKey>} Derived key
 */
async function deriveKeyFromPassword(password) {
  const encoder = new TextEncoder();
  const passwordBuffer = encoder.encode(password);

  // Import password as key material
  const passwordKey = await window.crypto.subtle.importKey(
    'raw',
    passwordBuffer,
    'PBKDF2',
    false,
    ['deriveKey']
  );

  // Get or generate salt
  let salt = localStorage.getItem('keySalt');
  if (!salt) {
    const saltBuffer = window.crypto.getRandomValues(new Uint8Array(16));
    salt = arrayBufferToBase64(saltBuffer);
    localStorage.setItem('keySalt', salt);
  }

  // Derive AES key
  return await window.crypto.subtle.deriveKey(
    {
      name: 'PBKDF2',
      salt: base64ToArrayBuffer(salt),
      iterations: 100000,
      hash: 'SHA-256'
    },
    passwordKey,
    {
      name: 'AES-GCM',
      length: 256
    },
    false,
    ['encrypt', 'decrypt']
  );
}

/**
 * Encrypt message content
 * @param {string} message - Plain text message
 * @param {string[]} recipientPublicKeys - Array of recipient public keys (base64)
 * @returns {Promise<Object>} Encrypted message bundle
 */
export async function encryptMessage(message, recipientPublicKeys) {
  // Generate random AES key for this message
  const messageKey = await window.crypto.subtle.generateKey(
    {
      name: 'AES-GCM',
      length: 256
    },
    true,
    ['encrypt', 'decrypt']
  );

  // Encrypt message with AES key
  const iv = window.crypto.getRandomValues(new Uint8Array(12));
  const encoder = new TextEncoder();
  const messageBuffer = encoder.encode(message);

  const encryptedMessage = await window.crypto.subtle.encrypt(
    {
      name: 'AES-GCM',
      iv: iv
    },
    messageKey,
    messageBuffer
  );

  // Export AES key
  const exportedMessageKey = await window.crypto.subtle.exportKey('raw', messageKey);

  // Encrypt AES key for each recipient with their RSA public key
  const encryptedKeys = {};
  const failedRecipients = [];

  for (const recipientKeyB64 of recipientPublicKeys) {
    try {
      // Import recipient's public key
      const recipientKey = await window.crypto.subtle.importKey(
        'spki',
        base64ToArrayBuffer(recipientKeyB64),
        {
          name: 'RSA-OAEP',
          hash: 'SHA-256'
        },
        true,
        ['encrypt']
      );

      // Encrypt message key with recipient's public key
      const encryptedKey = await window.crypto.subtle.encrypt(
        {
          name: 'RSA-OAEP'
        },
        recipientKey,
        exportedMessageKey
      );

      encryptedKeys[recipientKeyB64] = arrayBufferToBase64(encryptedKey);
    } catch (error) {
      failedRecipients.push({
        publicKey: recipientKeyB64.substring(0, 20) + '...', // Truncate for error message
        error: error.message
      });
    }
  }

  // Throw error if any encryptions failed to prevent partial message delivery
  if (failedRecipients.length > 0) {
    throw new Error(
      `Failed to encrypt message for ${failedRecipients.length} recipient(s): ${
        failedRecipients.map(r => r.error).join(', ')
      }`
    );
  }

  return {
    ciphertext: arrayBufferToBase64(encryptedMessage),
    iv: arrayBufferToBase64(iv),
    encryptedKeys: encryptedKeys // Map of publicKey -> encrypted AES key
  };
}

/**
 * Decrypt message content
 * @param {Object} encryptedBundle - Encrypted message bundle
 * @param {string} userPassword - User's password (to decrypt private key)
 * @param {string} userPublicKey - User's public key (to find correct encrypted key)
 * @returns {Promise<string>} Decrypted message
 */
export async function decryptMessage(encryptedBundle, userPassword, userPublicKey) {
  // Get user's private key
  const privateKey = await getPrivateKey(userPassword);

  // Find the encrypted key for this user
  const encryptedKeyB64 = encryptedBundle.encryptedKeys[userPublicKey];

  if (!encryptedKeyB64) {
    throw new Error('No encrypted key found for this user');
  }

  // Decrypt the AES key with user's private key
  const decryptedKeyBuffer = await window.crypto.subtle.decrypt(
    {
      name: 'RSA-OAEP'
    },
    privateKey,
    base64ToArrayBuffer(encryptedKeyB64)
  );

  // Import AES key
  const messageKey = await window.crypto.subtle.importKey(
    'raw',
    decryptedKeyBuffer,
    {
      name: 'AES-GCM',
      length: 256
    },
    false,
    ['decrypt']
  );

  // Decrypt message
  const decryptedBuffer = await window.crypto.subtle.decrypt(
    {
      name: 'AES-GCM',
      iv: base64ToArrayBuffer(encryptedBundle.iv)
    },
    messageKey,
    base64ToArrayBuffer(encryptedBundle.ciphertext)
  );

  // Convert to text
  const decoder = new TextDecoder();
  return decoder.decode(decryptedBuffer);
}

/**
 * Convert ArrayBuffer to Base64 string
 * @param {ArrayBuffer} buffer
 * @returns {string}
 */
function arrayBufferToBase64(buffer) {
  const bytes = new Uint8Array(buffer);
  let binary = '';
  for (let i = 0; i < bytes.byteLength; i++) {
    binary += String.fromCharCode(bytes[i]);
  }
  return window.btoa(binary);
}

/**
 * Convert Base64 string to ArrayBuffer
 * @param {string} base64
 * @returns {ArrayBuffer}
 */
function base64ToArrayBuffer(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;
}

/**
 * Check if encryption keys exist
 * @returns {boolean}
 */
export function hasEncryptionKeys() {
  return !!(
    localStorage.getItem('encryptedPrivateKey') &&
    localStorage.getItem('privateKeyIV') &&
    localStorage.getItem('keySalt')
  );
}

/**
 * Clear all stored encryption keys
 */
export function clearEncryptionKeys() {
  localStorage.removeItem('encryptedPrivateKey');
  localStorage.removeItem('privateKeyIV');
  localStorage.removeItem('keySalt');
}
