import {
  BaseBlock,
  Integer,
  Null,
  ObjectIdentifier,
  OctetString,
  Sequence,
  fromBER,
} from "asn1js";

const PBES2 = "1.2.840.113549.1.5.13";
const PBKDF2 = "1.2.840.113549.1.5.12";
const MAX_ENCRYPTED_BYTES = 8 * 1024 * 1024;
const MAX_ITERATIONS = 10_000_000;
export const MAX_TOTAL_PBKDF2_ITERATIONS = 20_000_000;

const PRFS: Record<string, string> = {
  "1.2.840.113549.2.7": "SHA-1",
  "1.2.840.113549.2.9": "SHA-256",
  "1.2.840.113549.2.10": "SHA-384",
  "1.2.840.113549.2.11": "SHA-512",
};

const CIPHERS: Record<
  string,
  { name: "AES-CBC" | "AES-GCM"; length: 128 | 192 | 256 }
> = {
  "2.16.840.1.101.3.4.1.2": { name: "AES-CBC", length: 128 },
  "2.16.840.1.101.3.4.1.22": { name: "AES-CBC", length: 192 },
  "2.16.840.1.101.3.4.1.42": { name: "AES-CBC", length: 256 },
  "2.16.840.1.101.3.4.1.6": { name: "AES-GCM", length: 128 },
  "2.16.840.1.101.3.4.1.26": { name: "AES-GCM", length: 192 },
  "2.16.840.1.101.3.4.1.46": { name: "AES-GCM", length: 256 },
};

const PRIVATE_KEY_ALGORITHMS: Record<string, string> = {
  "1.2.840.113549.1.1.1": "RSA",
  "1.2.840.113549.1.1.10": "RSA-PSS",
  "1.2.840.10045.2.1": "EC",
  "1.3.101.112": "Ed25519",
  "1.3.101.113": "Ed448",
  "1.3.101.110": "X25519",
  "1.3.101.111": "X448",
};

const CURVES: Record<string, string> = {
  "1.2.840.10045.3.1.7": "P-256",
  "1.3.132.0.34": "P-384",
  "1.3.132.0.35": "P-521",
};

export interface Pbkdf2WorkBudget {
  iterations: number;
}

export function reservePbkdf2Work(
  budget: Pbkdf2WorkBudget,
  iterations: number,
  label: string,
): void {
  if (
    !Number.isSafeInteger(budget.iterations) ||
    budget.iterations < 0 ||
    !Number.isSafeInteger(iterations) ||
    iterations < 1 ||
    budget.iterations > MAX_TOTAL_PBKDF2_ITERATIONS - iterations
  ) {
    throw new Error(
      `${label} exceeds the ${MAX_TOTAL_PBKDF2_ITERATIONS.toLocaleString()}-iteration aggregate PBKDF2 safety budget. Inspect fewer encrypted objects at once.`,
    );
  }
  budget.iterations += iterations;
}

type Block = BaseBlock;

function asBuffer(bytes: Uint8Array): ArrayBuffer {
  return bytes.buffer.slice(
    bytes.byteOffset,
    bytes.byteOffset + bytes.byteLength,
  ) as ArrayBuffer;
}

function parseDer(bytes: Uint8Array, name: string): Block {
  if (bytes.byteLength === 0 || bytes.byteLength > MAX_ENCRYPTED_BYTES)
    throw new Error(`${name} must contain 1 byte–8 MiB.`);
  const parsed = fromBER(asBuffer(bytes));
  if (parsed.offset === -1 || parsed.offset !== bytes.byteLength)
    throw new Error(`${name} is malformed DER or contains trailing bytes.`);
  const canonical = new Uint8Array(parsed.result.toBER(false));
  if (
    canonical.length !== bytes.length ||
    canonical.some((value, index) => value !== bytes[index])
  )
    throw new Error(`${name} must use canonical definite-length DER.`);
  return parsed.result;
}

function sequence(block: Block | undefined, name: string): Block[] {
  if (!(block instanceof Sequence))
    throw new Error(`${name} must be an ASN.1 SEQUENCE.`);
  return block.valueBlock.value;
}

function oid(block: Block | undefined, name: string): string {
  if (!(block instanceof ObjectIdentifier))
    throw new Error(`${name} must be an ASN.1 object identifier.`);
  return block.getValue();
}

function octets(block: Block | undefined, name: string): Uint8Array {
  if (!(block instanceof OctetString))
    throw new Error(`${name} must be an ASN.1 OCTET STRING.`);
  return new Uint8Array(block.getValue());
}

function integer(block: Block | undefined, name: string): bigint {
  if (!(block instanceof Integer))
    throw new Error(`${name} must be an ASN.1 INTEGER.`);
  return block.toBigInt();
}

export interface Pbes2Inspection {
  scheme: "PBES2";
  kdf: "PBKDF2";
  prf: string;
  iterations: number;
  saltBytes: number;
  cipher: "AES-CBC" | "AES-GCM";
  keyLength: 128 | 192 | 256;
  ivOrNonceBytes: number;
  tagLength?: number;
  encryptedBytes: number;
}

interface ParsedPbes2 extends Pbes2Inspection {
  salt: Uint8Array;
  ivOrNonce: Uint8Array;
  encryptedData: Uint8Array;
}

function parsePbes2(bytes: Uint8Array): ParsedPbes2 {
  const root = sequence(
    parseDer(bytes, "Encrypted PKCS #8"),
    "EncryptedPrivateKeyInfo",
  );
  if (root.length !== 2)
    throw new Error(
      "EncryptedPrivateKeyInfo must contain an algorithm and encrypted data.",
    );
  return parsePbes2Algorithm(
    root[0],
    octets(root[1], "Encrypted PKCS #8 data"),
  );
}

function parsePbes2Algorithm(
  algorithmBlock: Block | undefined,
  encryptedData: Uint8Array,
): ParsedPbes2 {
  const algorithm = sequence(algorithmBlock, "Encryption algorithm");
  if (
    algorithm.length !== 2 ||
    oid(algorithm[0], "Encryption algorithm") !== PBES2
  )
    throw new Error("Only PKCS #5 PBES2 encrypted PKCS #8 is supported.");
  const params = sequence(algorithm[1], "PBES2 parameters");
  if (params.length !== 2) throw new Error("PBES2 parameters are incomplete.");

  const kdf = sequence(params[0], "PBES2 key derivation function");
  if (kdf.length !== 2 || oid(kdf[0], "KDF algorithm") !== PBKDF2)
    throw new Error("Only PBKDF2 key derivation is supported for PBES2.");
  const pbkdf = sequence(kdf[1], "PBKDF2 parameters");
  if (pbkdf.length < 2 || pbkdf.length > 4)
    throw new Error("PBKDF2 parameters have an unsupported shape.");
  const salt = octets(pbkdf[0], "PBKDF2 salt");
  if (salt.length < 8 || salt.length > 1_024)
    throw new Error("PBKDF2 salt must contain 8–1,024 bytes.");
  const iterationBig = integer(pbkdf[1], "PBKDF2 iteration count");
  if (iterationBig < 1n || iterationBig > BigInt(MAX_ITERATIONS))
    throw new Error(
      `PBKDF2 iteration count must be 1–${MAX_ITERATIONS.toLocaleString()}.`,
    );
  let cursor = 2;
  let declaredKeyLength: number | undefined;
  if (pbkdf[cursor] instanceof Integer) {
    const bytesLong = integer(pbkdf[cursor], "PBKDF2 key length");
    if (bytesLong < 1n || bytesLong > 64n)
      throw new Error("PBKDF2 key length is outside the supported bound.");
    declaredKeyLength = Number(bytesLong) * 8;
    cursor += 1;
  }
  let prf = "SHA-1";
  if (pbkdf[cursor]) {
    const prfAlgorithm = sequence(pbkdf[cursor], "PBKDF2 PRF");
    if (
      prfAlgorithm.length < 1 ||
      prfAlgorithm.length > 2 ||
      (prfAlgorithm[1] !== undefined && !(prfAlgorithm[1] instanceof Null))
    )
      throw new Error("PBKDF2 PRF parameters must be absent or NULL.");
    prf = PRFS[oid(prfAlgorithm[0], "PBKDF2 PRF algorithm")] ?? "";
    if (!prf)
      throw new Error(
        "PBKDF2 PRF is not supported by this WebCrypto workflow.",
      );
    cursor += 1;
  }
  if (cursor !== pbkdf.length)
    throw new Error("PBKDF2 parameters contain unsupported trailing fields.");

  const encryption = sequence(params[1], "PBES2 encryption scheme");
  if (encryption.length !== 2)
    throw new Error("PBES2 encryption scheme parameters are incomplete.");
  const cipher = CIPHERS[oid(encryption[0], "PBES2 cipher")];
  if (!cipher)
    throw new Error("PBES2 cipher is not a supported AES-CBC/AES-GCM scheme.");
  if (declaredKeyLength !== undefined && declaredKeyLength !== cipher.length)
    throw new Error(
      "PBKDF2 declared key length does not match the AES scheme.",
    );

  let ivOrNonce: Uint8Array;
  let tagLength: number | undefined;
  if (cipher.name === "AES-CBC") {
    ivOrNonce = octets(encryption[1], "AES-CBC IV");
    if (ivOrNonce.length !== 16)
      throw new Error("AES-CBC IV must contain 16 bytes.");
  } else {
    const gcm = sequence(encryption[1], "AES-GCM parameters");
    if (gcm.length < 1 || gcm.length > 2)
      throw new Error("AES-GCM parameters have an unsupported shape.");
    ivOrNonce = octets(gcm[0], "AES-GCM nonce");
    if (ivOrNonce.length < 12 || ivOrNonce.length > 16)
      throw new Error("AES-GCM nonce must contain 12–16 bytes.");
    const tagBytes = gcm[1] ? integer(gcm[1], "AES-GCM tag length") : 12n;
    if (![12n, 13n, 14n, 15n, 16n].includes(tagBytes))
      throw new Error("AES-GCM tag length must contain 12–16 octets.");
    tagLength = Number(tagBytes) * 8;
  }
  if (encryptedData.length === 0 || encryptedData.length > MAX_ENCRYPTED_BYTES)
    throw new Error("PBES2 encrypted payload must contain 1 byte–8 MiB.");
  return {
    scheme: "PBES2",
    kdf: "PBKDF2",
    prf,
    iterations: Number(iterationBig),
    saltBytes: salt.length,
    cipher: cipher.name,
    keyLength: cipher.length,
    ivOrNonceBytes: ivOrNonce.length,
    ...(tagLength ? { tagLength } : {}),
    encryptedBytes: encryptedData.length,
    salt,
    ivOrNonce,
    encryptedData,
  };
}

export function inspectEncryptedPkcs8(bytes: Uint8Array): Pbes2Inspection {
  const parsed = parsePbes2(bytes);
  return {
    scheme: parsed.scheme,
    kdf: parsed.kdf,
    prf: parsed.prf,
    iterations: parsed.iterations,
    saltBytes: parsed.saltBytes,
    cipher: parsed.cipher,
    keyLength: parsed.keyLength,
    ivOrNonceBytes: parsed.ivOrNonceBytes,
    ...(parsed.tagLength === undefined ? {} : { tagLength: parsed.tagLength }),
    encryptedBytes: parsed.encryptedBytes,
  };
}

export function inspectPbes2Payload(
  algorithmIdentifier: Uint8Array,
  encryptedData: Uint8Array,
): Pbes2Inspection {
  const parsed = parsePbes2Algorithm(
    parseDer(algorithmIdentifier, "PBES2 AlgorithmIdentifier"),
    encryptedData,
  );
  return publicPbes2Inspection(parsed);
}

export interface Pkcs8Inspection {
  algorithmOid: string;
  algorithm: string;
  curve?: string;
  bytes: number;
}

export function inspectPkcs8(bytes: Uint8Array): Pkcs8Inspection {
  const root = sequence(
    parseDer(bytes, "PKCS #8 private key"),
    "PrivateKeyInfo",
  );
  if (root.length < 3) throw new Error("PKCS #8 PrivateKeyInfo is incomplete.");
  const version = integer(root[0], "PKCS #8 version");
  if (version < 0n || version > 1n)
    throw new Error("PKCS #8 version is unsupported.");
  const algorithmIdentifier = sequence(root[1], "Private-key algorithm");
  const algorithmOid = oid(algorithmIdentifier[0], "Private-key algorithm");
  const algorithm = PRIVATE_KEY_ALGORITHMS[algorithmOid] ?? "Unknown";
  const curveOid =
    algorithm === "EC" && algorithmIdentifier[1] instanceof ObjectIdentifier
      ? algorithmIdentifier[1].getValue()
      : undefined;
  // Validates that the privateKey field is present and encoded as octets.
  octets(root[2], "PKCS #8 privateKey");
  return {
    algorithmOid,
    algorithm,
    ...(curveOid ? { curve: CURVES[curveOid] ?? `OID ${curveOid}` } : {}),
    bytes: bytes.length,
  };
}

export async function decryptEncryptedPkcs8(
  bytes: Uint8Array,
  password: string,
): Promise<{
  bytes: Uint8Array;
  inspection: Pbes2Inspection;
  key: Pkcs8Inspection;
}> {
  if (password.length > 100_000)
    throw new Error("Password must contain at most 100,000 UTF-16 units.");
  const parsed = parsePbes2(bytes);
  const output = await decryptParsedPbes2(parsed, password);
  try {
    return {
      bytes: output,
      inspection: publicPbes2Inspection(parsed),
      key: inspectPkcs8(output),
    };
  } catch (error) {
    output.fill(0);
    throw error;
  }
}

export async function decryptPbes2Payload(
  algorithmIdentifier: Uint8Array,
  encryptedData: Uint8Array,
  password: string,
): Promise<{ bytes: Uint8Array; inspection: Pbes2Inspection }> {
  if (password.length > 100_000)
    throw new Error("Password must contain at most 100,000 UTF-16 units.");
  const parsed = parsePbes2Algorithm(
    parseDer(algorithmIdentifier, "PBES2 AlgorithmIdentifier"),
    encryptedData,
  );
  return {
    bytes: await decryptParsedPbes2(parsed, password),
    inspection: publicPbes2Inspection(parsed),
  };
}

function publicPbes2Inspection(parsed: ParsedPbes2): Pbes2Inspection {
  return {
    scheme: parsed.scheme,
    kdf: parsed.kdf,
    prf: parsed.prf,
    iterations: parsed.iterations,
    saltBytes: parsed.saltBytes,
    cipher: parsed.cipher,
    keyLength: parsed.keyLength,
    ivOrNonceBytes: parsed.ivOrNonceBytes,
    ...(parsed.tagLength === undefined ? {} : { tagLength: parsed.tagLength }),
    encryptedBytes: parsed.encryptedBytes,
  };
}

async function decryptParsedPbes2(
  parsed: ParsedPbes2,
  password: string,
): Promise<Uint8Array> {
  const material = await crypto.subtle.importKey(
    "raw",
    new TextEncoder().encode(password),
    "PBKDF2",
    false,
    ["deriveKey"],
  );
  const key = await crypto.subtle.deriveKey(
    {
      name: "PBKDF2",
      salt: asBuffer(parsed.salt),
      iterations: parsed.iterations,
      hash: parsed.prf,
    },
    material,
    { name: parsed.cipher, length: parsed.keyLength },
    false,
    ["decrypt"],
  );
  let decrypted: ArrayBuffer;
  try {
    decrypted = await crypto.subtle.decrypt(
      parsed.cipher === "AES-CBC"
        ? { name: "AES-CBC", iv: asBuffer(parsed.ivOrNonce) }
        : {
            name: "AES-GCM",
            iv: asBuffer(parsed.ivOrNonce),
            tagLength: parsed.tagLength ?? 96,
          },
      key,
      asBuffer(parsed.encryptedData),
    );
  } catch {
    throw new Error(
      "PBES2 decryption failed. The password may be wrong, the data may be damaged, or this WebCrypto runtime may not support the parameters.",
    );
  }
  return new Uint8Array(decrypted);
}