import assert from "node:assert/strict";
import { test } from "node:test";
import {
  createRecipientKeyPair,
  decryptPayloadFromEnvelope,
  encryptBytesForRecipient,
  encryptBytesForRecipients,
  getEnvelopeRecipientWraps,
  rewrapEnvelopeForRecipients
} from "../dist/index.js";

test("multi-recipient envelope decrypts with each recipient private key", async () => {
  const first = await createRecipientKeyPair();
  const second = await createRecipientKeyPair();
  const plaintext = new TextEncoder().encode("shared encrypted payload");

  const encrypted = await encryptBytesForRecipients(plaintext.buffer as ArrayBuffer, [first.publicKey, second.publicKey], {
    purpose: "deliverable"
  });

  const wraps = getEnvelopeRecipientWraps(encrypted.envelope);
  assert.equal(wraps.length, 2);
  assert.deepEqual(
    wraps.map((wrap) => wrap.keyFingerprint).sort(),
    [first.fingerprint, second.fingerprint].sort()
  );

  const ciphertext = await encrypted.ciphertext.arrayBuffer();
  const firstDecrypted = await decryptPayloadFromEnvelope(ciphertext, encrypted.envelope, first.privateKey);
  const secondDecrypted = await decryptPayloadFromEnvelope(ciphertext, encrypted.envelope, second.privateKey);

  assert.equal(new TextDecoder().decode(firstDecrypted.plaintext), "shared encrypted payload");
  assert.equal(new TextDecoder().decode(secondDecrypted.plaintext), "shared encrypted payload");
  assert.deepEqual(firstDecrypted.metadata, { purpose: "deliverable" });
  assert.deepEqual(secondDecrypted.metadata, { purpose: "deliverable" });
});

test("single-recipient helper keeps the legacy recipient field", async () => {
  const recipient = await createRecipientKeyPair();
  const plaintext = new TextEncoder().encode("legacy payload");

  const encrypted = await encryptBytesForRecipient(plaintext.buffer as ArrayBuffer, recipient.publicKey, {
    purpose: "source"
  });

  assert.equal(encrypted.envelope.recipient.keyFingerprint, recipient.fingerprint);
  assert.equal(getEnvelopeRecipientWraps(encrypted.envelope).length, 1);

  const decrypted = await decryptPayloadFromEnvelope(await encrypted.ciphertext.arrayBuffer(), encrypted.envelope, recipient.privateKey);
  assert.equal(new TextDecoder().decode(decrypted.plaintext), "legacy payload");
  assert.deepEqual(decrypted.metadata, { purpose: "source" });
});

test("envelope can be rewrapped to a new recipient without changing ciphertext metadata", async () => {
  const admin = await createRecipientKeyPair();
  const client = await createRecipientKeyPair();
  const replacementClient = await createRecipientKeyPair();
  const plaintext = new TextEncoder().encode("deliverable payload");

  const encrypted = await encryptBytesForRecipients(plaintext.buffer as ArrayBuffer, [admin.publicKey, client.publicKey], {
    purpose: "deliverable"
  });
  const originalContent = encrypted.envelope.content;
  const originalMetadata = encrypted.envelope.metadata;

  const rewrappedEnvelope = await rewrapEnvelopeForRecipients(
    encrypted.envelope,
    admin.privateKey,
    [replacementClient.publicKey],
    { keyFingerprint: admin.fingerprint }
  );

  assert.deepEqual(rewrappedEnvelope.content, originalContent);
  assert.deepEqual(rewrappedEnvelope.metadata, originalMetadata);
  assert.deepEqual(
    getEnvelopeRecipientWraps(rewrappedEnvelope)
      .map((wrap) => wrap.keyFingerprint)
      .sort(),
    [admin.fingerprint, client.fingerprint, replacementClient.fingerprint].sort()
  );

  const decrypted = await decryptPayloadFromEnvelope(
    await encrypted.ciphertext.arrayBuffer(),
    rewrappedEnvelope,
    replacementClient.privateKey
  );
  assert.equal(new TextDecoder().decode(decrypted.plaintext), "deliverable payload");
  assert.deepEqual(decrypted.metadata, { purpose: "deliverable" });
});