import { Buffer } from 'buffer';
import { sha256 } from "js-sha256";
import JSChaCha20 from 'js-chacha20';

function bigIntToUint8Array(bigInt, byteLength) {
  const hex = bigInt.toString(16).padStart(byteLength * 2, '0');
  const byteArray = new Uint8Array(byteLength);
  for (let i = 0; i < byteLength; i++) {
    byteArray[i] = parseInt(hex.substr(i * 2, 2), 16);
  }
  return byteArray;
}

export function bigIntToBase64(bigInt) {
  const hexString = bigInt.toString(16).padStart(bigInt.toString(16).length + (bigInt.toString(16).length % 2), '0');
  return Buffer.from(hexString, 'hex').toString('base64');
}

export function base64ToBigInt(base64) {
  return BigInt('0x' + Buffer.from(base64, 'base64').toString('hex'));
}

const pBase64 = "2IGKThw2+H+X0Mc5ZxIfOGX1b3lLc/mDB2Ne73FdoEiJRRfWaETlLUg8dAFUWn4Jxg/QSbP7+f/Q0dZbPCShAXrqWsqxScNk+XvFRTUAhqq1h82Bh4Puok2P1ke6sPR/k+LCt9uHMlt/X0TodDonFIMr87KjB9bQ+zysfPXU2G9chMjqPpY2AkInPHfaMtKtIfBslXKbhGwtaK6t0h0GGJ7W3GU8e1dzo/JgwVDnqoTryAKdoFmrpjg41naQOC5Hl59i+Ik5yEL+NCvSis7IYo55sM6cbu8B4N4wNwDKkefgElADvhSKQJirbmSpPXs5Lr7GXgRBR6t/AGrYTgahPw==";
const gBase64 = "5";
const p = base64ToBigInt(pBase64);
const g = BigInt(gBase64);

export function generatePrivateKey() {
  return generateRandomBigInt(p - 1n);
}

function generateRandomBigInt(max) {
  let randomHex = '';
  while (randomHex.length < max.toString(16).length) {
    randomHex += Math.floor(Math.random() * 16).toString(16);
  }
  return BigInt('0x' + randomHex.slice(0, max.toString(16).length));
}

export function calculatePublicKey(privateKey) {
  return modExponentiation(g, privateKey, p);
}

function modExponentiation(base, exponent, modulus) {
  let result = BigInt(1);
  base = base % modulus;
  while (exponent > 0n) {
    if (exponent % 2n === 1n) result = (result * base) % modulus;
    exponent /= 2n;
    base = (base * base) % modulus;
  }
  return result;
}

export function computeSharedSecret(publicKey, privateKey) {
  return modExponentiation(publicKey, privateKey, p);
}

export function chacha20Encrypt(jsonobj) {
  try {
    let { message, receiver, pubkey } = JSON.parse(jsonobj);
    message = new TextEncoder().encode(message);

    let storedUsers = JSON.parse(localStorage.getItem("privateKeys")) || [];
    const user = storedUsers.find(user => user.username === receiver);
    if (!user) return console.error("User not found in localStorage");

    const clientPrivKey = base64ToBigInt(user.private64);
    const secret = computeSharedSecret(base64ToBigInt(pubkey), clientPrivKey);

    const secretArray = bigIntToUint8Array(secret);
    const chachaKeyUint8 = new Uint8Array(sha256.create().update(secretArray).digest());

    const nonce = generateNonce();
    let nonceB64 = Buffer.from(nonce).toString("base64");

    const chacha = new JSChaCha20(chachaKeyUint8, nonce);
    const encrypted = chacha.encrypt(message);
    const encryptedArray = new Uint8Array(encrypted);
    const encryptedB64 = Buffer.from(encryptedArray).toString("base64");

    return JSON.stringify({
      encrypted: encryptedB64,
      nonce: nonceB64,
      pubKey: pubkey,
      receiverUsername: receiver
    });
  } catch (error) {
    console.error("Error during encryption:", error);
    return null;
  }
}

export function chacha20Decrypt(data) {
  try {
    const { content, nonce, pubKey, receiverUsername } = data;

    let nonceArray = new Uint8Array(Buffer.from(nonce, 'base64'));
    if (nonceArray.length !== 12) throw new Error("Nonce should be a 12 byte array!");

    let storedUsers = JSON.parse(localStorage.getItem("privateKeys")) || [];
    const user = storedUsers.find(user => user.username === receiverUsername);
    if (!user) throw new Error("User not found in localStorage for decryption");

    const clientPrivKey = base64ToBigInt(user.private64);
    const secret = computeSharedSecret(base64ToBigInt(pubKey), clientPrivKey);

    const secretArray = bigIntToUint8Array(secret);
    const chachaKeyUint8 = new Uint8Array(sha256.create().update(secretArray).digest());

    let encryptedArray = new Uint8Array(Buffer.from(content, 'base64'));

    const chacha = new JSChaCha20(chachaKeyUint8, nonceArray);
    const decrypted = chacha.decrypt(encryptedArray);

    return new TextDecoder().decode(decrypted);
  } catch (error) {
    console.error("Decryption Error:", error);
    return "Error decrypting message: " + error.message;
  }
}

export function generateNonce(length = 12) {
  const nonce = new Uint8Array(length);
  for (let i = 0; i < length; i++) nonce[i] = Math.floor(Math.random() * 256);
  return nonce;
}