import * as forge from "node-forge";

import { Utils } from "@bitwarden/common/platform/misc/utils";

import { RecordKeyType } from "../generated/record_pb";
import {
  KeeperKey,
  KeeperNonce,
  KeeperSalt,
  RsaPrivateKey,
  RsaPublicKey,
} from "../models/crypto-types";

const AES_GCM_NONCE_SIZE = 12;
const AES_GCM_TAG_SIZE = 16;
const AES_BLOCK_SIZE = 16;

export function getRandomBytes(length: number): Uint8Array {
  const arr = new Uint8Array(length);
  crypto.getRandomValues(arr);
  return arr;
}

export function generateEncryptionKey(): KeeperKey {
  return getRandomBytes(32) as KeeperKey;
}

/**
 * Decrypt an aes-v1 packet. Keeper aes-v1 is AES-CBC without any authentication (MAC).
 */
export async function decryptAesV1(data: Uint8Array, key: KeeperKey): Promise<Uint8Array> {
  const iv = data.subarray(0, AES_BLOCK_SIZE);
  const encrypted = data.subarray(AES_BLOCK_SIZE);
  const cryptoKey = await crypto.subtle.importKey(
    "raw",
    toBuffer(key),
    { name: "AES-CBC" },
    false,
    ["decrypt"],
  );
  const decrypted = await crypto.subtle.decrypt(
    { name: "AES-CBC", iv: toBuffer(iv) },
    cryptoKey,
    toBuffer(encrypted),
  );
  return new Uint8Array(decrypted);
}

/**
 * Encrypt a packet with aes-v2. Keeper aes-v2 is AES-GCM. Optionally accepts an external nonce,
 * if not provided the nonce will be generated automatically.
 * WARNING: DO NOT REUSE THE NONCE. THIS WILL LEAD TO CATASTROPHIC CRYPTOGRAPHIC FAILURE.
 */
export async function encryptAesV2(
  data: Uint8Array,
  key: KeeperKey,
  nonce?: KeeperNonce,
  nonceLength = AES_GCM_NONCE_SIZE,
): Promise<Uint8Array> {
  const nonceBuffer = nonce ?? (getRandomBytes(nonceLength) as KeeperNonce);
  const cryptoKey = await crypto.subtle.importKey(
    "raw",
    toBuffer(key),
    { name: "AES-GCM" },
    false,
    ["encrypt"],
  );
  const encrypted = await crypto.subtle.encrypt(
    { name: "AES-GCM", iv: toBuffer(nonceBuffer), tagLength: AES_GCM_TAG_SIZE * 8 },
    cryptoKey,
    toBuffer(data),
  );
  return concatUint8Arrays(nonceBuffer, new Uint8Array(encrypted));
}

/**
 * Decrypt a packet with aes-v2. Keeper aes-v2 is AES-GCM. The packet has the format:
 * Nonce|Plaintext|Tag where Plaintext|Tag are the output of AES-GCM.
 */
export async function decryptAesV2(
  data: Uint8Array,
  key: KeeperKey,
  nonceLength = AES_GCM_NONCE_SIZE,
): Promise<Uint8Array> {
  const nonce = data.subarray(0, nonceLength);
  // Web Crypto API expects ciphertext + tag together (not separated like Node.js)
  const ciphertextWithTag = data.subarray(nonceLength);
  const cryptoKey = await crypto.subtle.importKey(
    "raw",
    toBuffer(key),
    { name: "AES-GCM" },
    false,
    ["decrypt"],
  );
  const decrypted = await crypto.subtle.decrypt(
    { name: "AES-GCM", iv: toBuffer(nonce), tagLength: AES_GCM_TAG_SIZE * 8 },
    cryptoKey,
    toBuffer(ciphertextWithTag),
  );
  return new Uint8Array(decrypted);
}

export function encryptRsa(data: Uint8Array, publicKeyBytes: RsaPublicKey): Uint8Array {
  // Use node-forge for RSA PKCS#1 v1.5 padding (required by Keeper)
  // Web Crypto API doesn't support PKCS#1 v1.5 for encryption
  const publicKeyDer = forge.util.createBuffer(uint8ArrayToByteString(publicKeyBytes));
  const asn1 = forge.asn1.fromDer(publicKeyDer);
  const publicKey = forge.pki.publicKeyFromAsn1(asn1);

  const dataBytes = uint8ArrayToByteString(data);
  const encrypted = publicKey.encrypt(dataBytes, "RSAES-PKCS1-V1_5");
  return byteStringToUint8Array(encrypted);
}

export function decryptRsa(data: Uint8Array, privateKeyBytes: RsaPrivateKey): Uint8Array {
  // Use node-forge for RSA PKCS#1 v1.5 padding
  const privateKeyDer = forge.util.createBuffer(uint8ArrayToByteString(privateKeyBytes));
  const asn1 = forge.asn1.fromDer(privateKeyDer);
  const privateKey = forge.pki.privateKeyFromAsn1(asn1);

  const dataBytes = uint8ArrayToByteString(data);
  const decrypted = privateKey.decrypt(dataBytes, "RSAES-PKCS1-V1_5");
  return byteStringToUint8Array(decrypted);
}

export async function generateEcKey(): Promise<{ privateKey: CryptoKey; publicKey: CryptoKey }> {
  const keyPair = await crypto.subtle.generateKey({ name: "ECDH", namedCurve: "P-256" }, true, [
    "deriveBits",
  ]);
  return { privateKey: keyPair.privateKey, publicKey: keyPair.publicKey };
}

export async function unloadEcPublicKey(publicKey: CryptoKey): Promise<Uint8Array> {
  const exported = await crypto.subtle.exportKey("raw", publicKey);
  return new Uint8Array(exported);
}

export async function loadEcPublicKey(publicKeyBytes: Uint8Array): Promise<CryptoKey> {
  if (publicKeyBytes.length < 65 || publicKeyBytes[0] !== 0x04) {
    throw new Error("Invalid EC public key data");
  }
  return await crypto.subtle.importKey(
    "raw",
    toBuffer(publicKeyBytes),
    { name: "ECDH", namedCurve: "P-256" },
    true,
    [],
  );
}

export async function loadEcPrivateKey(privateKeyBytes: Uint8Array): Promise<CryptoKey> {
  return await crypto.subtle.importKey(
    "pkcs8",
    toBuffer(privateKeyBytes),
    { name: "ECDH", namedCurve: "P-256" },
    true,
    ["deriveBits"],
  );
}

async function deriveEcdhKey(publicKey: CryptoKey, privateKey: CryptoKey): Promise<KeeperKey> {
  const sharedSecret = await crypto.subtle.deriveBits(
    { name: "ECDH", public: publicKey },
    privateKey,
    256,
  );
  const hash = await crypto.subtle.digest("SHA-256", sharedSecret);
  return new Uint8Array(hash) as KeeperKey;
}

/**
 * Encrypts an ECC-encrypted packet. Keeper ECC encryption works by deriving a shared key content-encrytpion-key using ECDH, then encrypting the content with under the content-encryption-key with AES-GCM
 */
export async function encryptEc(data: Uint8Array, publicKey: CryptoKey): Promise<Uint8Array> {
  const { privateKey: ephemeralPrivate, publicKey: ephemeralPublic } = await generateEcKey();
  const encryptionKey = await deriveEcdhKey(publicKey, ephemeralPrivate);
  const encryptedData = await encryptAesV2(data, encryptionKey);
  const ephemeralPublicBytes = await unloadEcPublicKey(ephemeralPublic);
  return concatUint8Arrays(ephemeralPublicBytes, encryptedData);
}

/**
 * Decrypts an ECC-encrypted-packet. Keeper ECC encryption works by deriving an ephemeral content-encryption-key, then encrypting the payload using AES-GCM. The packet concatinates the ECC ephemeral public key, and the AES-GCM encrypted packet.
 */
export async function decryptEc(data: Uint8Array, privateKey: CryptoKey): Promise<Uint8Array> {
  const ephemeralPublicBytes = data.subarray(0, 65);
  const ephemeralPublic = await loadEcPublicKey(ephemeralPublicBytes);
  const encryptionKey = await deriveEcdhKey(ephemeralPublic, privateKey);
  const encryptedData = data.subarray(65);
  return await decryptAesV2(encryptedData, encryptionKey);
}

/**
 * Derives a keeper master-key using a password, salt and a number of iterations.
 */
export async function deriveKeyV1(
  password: string,
  salt: KeeperSalt,
  iterations: number,
): Promise<KeeperKey> {
  const passwordKey = await crypto.subtle.importKey(
    "raw",
    new TextEncoder().encode(password),
    "PBKDF2",
    false,
    ["deriveBits"],
  );
  const derivedBits = await crypto.subtle.deriveBits(
    {
      name: "PBKDF2",
      salt: toBuffer(salt),
      iterations,
      hash: "SHA-256",
    },
    passwordKey,
    256,
  );
  return new Uint8Array(derivedBits) as KeeperKey;
}

export async function deriveV1KeyHash(
  password: string,
  salt: KeeperSalt,
  iterations: number,
): Promise<Uint8Array> {
  const key = await deriveKeyV1(password, salt, iterations);
  const hash = await crypto.subtle.digest("SHA-256", toBuffer(key));
  return new Uint8Array(hash);
}

/**
 * Derives a key from the encryption parameters. The encryption parameters is a binary blob that
 * contains the data needed to derive the key. It is structured as follows.
 * CORRUPTED[1]|Iterations[3]|Salt[16]|IV[16]|Data
 * The iterations, salt, key are used to derive a content-encryption-key. The IV and cek are used
 * to decrypt the data, which consists of two 32-byte blocks. These two blocks need to match.
 * They are decrypted with AES-CBC, using no pkcs7 padding.
 */
export async function decryptEncryptionParams(
  password: string,
  encryptionParams: Uint8Array,
): Promise<KeeperKey> {
  const CORRUPTED_MESSAGE = "Corrupted encryption parameters";

  if (encryptionParams[0] !== 1) {
    throw new Error(CORRUPTED_MESSAGE);
  }

  if (encryptionParams.length !== 1 + 3 + 16 + 16 + 64) {
    throw new Error(CORRUPTED_MESSAGE);
  }

  const iterations = (encryptionParams[1] << 16) + (encryptionParams[2] << 8) + encryptionParams[3];
  const salt = encryptionParams.subarray(4, 20) as KeeperSalt;
  const key = await deriveKeyV1(password, salt, iterations);

  // Decrypt data with no padding
  // We need to manually handle this since Web Crypto always expects PKCS7 padding
  const iv = encryptionParams.subarray(20, 36) as KeeperNonce;
  const encryptedData = encryptionParams.subarray(36);
  const decrypted = await decryptAesNoPadding(encryptedData, key, iv);

  const first32 = decrypted.subarray(0, 32);
  const second32 = decrypted.subarray(32, 64);
  if (!uint8ArrayEquals(first32, second32)) {
    throw new Error(CORRUPTED_MESSAGE);
  }

  return first32 as KeeperKey;
}

async function decryptAesNoPadding(
  data: Uint8Array,
  key: KeeperKey,
  iv: KeeperNonce,
): Promise<Uint8Array> {
  // Web Crypto doesn't support no-padding mode, so we need a workaround.
  // AES-CBC chains individual aes encryption operation blocks using the previous ciphertext
  // as the IV of the next block. In PKCS7, the last block needs to be valid padding. In the
  // case that the full block is padding and contains no data, the entire block should contain
  // 0x10. Because of this, a non-padded AES-CBC block can be converted to a correctly padded
  // block, in the case that the plaintext length is a multiple of the block size. This is done
  // by using the last ciphertext block as the IV to encrypt the padding block containing only
  // 0x10. Appending this at the end, then results in a valid PKCS7-padded block.
  const cryptoKey = await crypto.subtle.importKey(
    "raw",
    toBuffer(key),
    { name: "AES-CBC" },
    false,
    ["encrypt", "decrypt"],
  );

  // Get the last block of ciphertext to use as IV for encrypting the padding
  const lastBlock = data.subarray(data.length - AES_BLOCK_SIZE);

  // Create a proper PKCS7 padding block (16 bytes of 0x10)
  const paddingBlock = new Uint8Array(AES_BLOCK_SIZE).fill(AES_BLOCK_SIZE);

  // Encrypt the padding block using the last ciphertext block as IV
  const encryptedPadding = await crypto.subtle.encrypt(
    { name: "AES-CBC", iv: toBuffer(lastBlock) },
    cryptoKey,
    paddingBlock,
  );

  // Take only the first block of the encrypted padding (the second block is the padding of the padding)
  const encryptedPaddingBlock = new Uint8Array(encryptedPadding).subarray(0, AES_BLOCK_SIZE);

  // Append the encrypted padding to the original ciphertext
  const paddedCiphertext = concatUint8Arrays(data, encryptedPaddingBlock);

  // Now decrypt - Web Crypto will find valid PKCS7 padding
  const decrypted = await crypto.subtle.decrypt(
    { name: "AES-CBC", iv: toBuffer(iv) },
    cryptoKey,
    toBuffer(paddedCiphertext),
  );

  // Return only the original data length (strip the decrypted padding block)
  return new Uint8Array(decrypted).subarray(0, data.length);
}

export async function decryptKeeperKey(
  encryptedKey: Uint8Array,
  keyType: RecordKeyType,
  dataKey: KeeperKey,
  rsaPrivateKey?: RsaPrivateKey,
  ecPrivateKey?: CryptoKey,
): Promise<KeeperKey> {
  switch (keyType) {
    case RecordKeyType.NO_KEY:
      throw new Error("Cannot decrypt NO_KEY type");

    case RecordKeyType.ENCRYPTED_BY_DATA_KEY:
      return (await decryptAesV1(encryptedKey, dataKey)) as KeeperKey;

    case RecordKeyType.ENCRYPTED_BY_PUBLIC_KEY:
      if (!rsaPrivateKey) {
        throw new Error("RSA private key required for ENCRYPTED_BY_PUBLIC_KEY");
      }
      return decryptRsa(encryptedKey, rsaPrivateKey) as KeeperKey;

    case RecordKeyType.ENCRYPTED_BY_DATA_KEY_GCM:
      return (await decryptAesV2(encryptedKey, dataKey)) as KeeperKey;

    case RecordKeyType.ENCRYPTED_BY_PUBLIC_KEY_ECC:
      if (!ecPrivateKey) {
        throw new Error("EC private key required for ENCRYPTED_BY_PUBLIC_KEY_ECC");
      }
      return (await decryptEc(encryptedKey, ecPrivateKey)) as KeeperKey;

    default:
      throw new Error(`Unknown key type: ${keyType}`);
  }
}

export function base64UrlEncode(data: Uint8Array): string {
  return Utils.fromB64toUrlB64(Utils.fromBufferToB64(data));
}

export function base64UrlDecode(text: string): Uint8Array {
  return Utils.fromUrlB64ToArray(text);
}

// Helper functions

function concatUint8Arrays(...arrays: Uint8Array[]): Uint8Array {
  const totalLength = arrays.reduce((sum, arr) => sum + arr.length, 0);
  const result = new Uint8Array(totalLength);
  let offset = 0;
  for (const arr of arrays) {
    result.set(arr, offset);
    offset += arr.length;
  }
  return result;
}

// Constant-time equality: always compares every byte (no early return on the first mismatch) so
// the running time does not depend on where the inputs first differ. The length check is not
// timing-sensitive here — the inputs are fixed-size blocks.
function uint8ArrayEquals(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) {
    return false;
  }
  let diff = 0;
  for (let i = 0; i < a.length; i++) {
    diff |= a[i] ^ b[i];
  }
  return diff === 0;
}

function uint8ArrayToByteString(data: Uint8Array): string {
  let result = "";
  for (let i = 0; i < data.length; i++) {
    result += String.fromCharCode(data[i]);
  }
  return result;
}

function toBuffer(data: Uint8Array): BufferSource {
  return data as unknown as BufferSource;
}

function byteStringToUint8Array(str: string): Uint8Array {
  const result = new Uint8Array(str.length);
  for (let i = 0; i < str.length; i++) {
    result[i] = str.charCodeAt(i);
  }
  return result;
}
