// Run: npm test   (Node strips the TypeScript types itself)
// Pure wire-format tests: no network, no IndexedDB — the reason omemo-wire
// is its own module, separate from omemo.ts.
import { test } from "node:test";
import assert from "node:assert/strict";
import parse from "ltx/lib/parse.js";
import { toBase64 } from "../src/xmpp/push.ts";
import { buildBundle, buildDeviceList, buildEncryptedMessage, parseBundle, parseDeviceList, parseEncryptedMessage } from "../src/crypto/omemo-wire.ts";

const NS = "eu.siacs.conversations.axolotl";
const el = (s: string) => parse(s) as any;

// Distinct key material: every byte pattern different, so a mix-up between
// the parts fails the round-trip loudly.
const key = (byte: number, length = 32) => new Uint8Array(length).fill(byte).buffer;

test("bundle round-trips through build→parse, in the legacy wire format", () => {
  const identityKey = key(1, 33);
  const signedPreKey = { keyId: 7, publicKey: key(2), signature: key(3, 64) };
  const preKeys = [
    { keyId: 1, publicKey: key(4) },
    { keyId: 42, publicKey: key(5) },
  ];
  const built = buildBundle(identityKey, signedPreKey, preKeys);
  // Lock the exact wire layout against XEP-0384 (legacy), not just against
  // our own parser: a round-trip through our own build+parse pair would
  // happily pass if both drifted together away from what Monal reads.
  assert.equal(
    built.toString(),
    `<bundle xmlns="${NS}">` +
      `<signedPreKeyPublic signedPreKeyId="7">${toBase64(signedPreKey.publicKey)}</signedPreKeyPublic>` +
      `<signedPreKeySignature>${toBase64(signedPreKey.signature)}</signedPreKeySignature>` +
      `<identityKey>${toBase64(identityKey)}</identityKey>` +
      `<prekeys>` +
      `<preKeyPublic preKeyId="1">${toBase64(preKeys[0].publicKey)}</preKeyPublic>` +
      `<preKeyPublic preKeyId="42">${toBase64(preKeys[1].publicKey)}</preKeyPublic>` +
      `</prekeys>` +
      `</bundle>`,
  );
  const parsed = parseBundle(built)!;
  assert.deepEqual(new Uint8Array(parsed.identityKey), new Uint8Array(identityKey));
  assert.equal(parsed.signedPreKey.keyId, 7);
  assert.deepEqual(new Uint8Array(parsed.signedPreKey.publicKey), new Uint8Array(signedPreKey.publicKey));
  assert.deepEqual(new Uint8Array(parsed.signedPreKey.signature), new Uint8Array(signedPreKey.signature));
  assert.deepEqual(
    parsed.preKeys.map((k) => k.keyId),
    [1, 42],
  );
  assert.deepEqual(new Uint8Array(parsed.preKeys[1].publicKey), new Uint8Array(preKeys[1].publicKey));
});

test("a bundle without one-time prekeys round-trips (all consumed)", () => {
  const built = buildBundle(key(1, 33), { keyId: 1, publicKey: key(2), signature: key(3, 64) }, []);
  assert.deepEqual(parseBundle(built)!.preKeys, []);
  // The prekeys container itself may be absent, not just empty.
  const absent = el(
    `<bundle xmlns="${NS}">` +
      `<signedPreKeyPublic signedPreKeyId="1">${toBase64(key(2))}</signedPreKeyPublic>` +
      `<signedPreKeySignature>${toBase64(key(3, 64))}</signedPreKeySignature>` +
      `<identityKey>${toBase64(key(1, 33))}</identityKey>` +
      `</bundle>`,
  );
  assert.deepEqual(parseBundle(absent)!.preKeys, []);
});

test("a malformed or missing bundle parses to null, never throws", () => {
  const B = (inner: string) => el(`<bundle xmlns="${NS}">${inner}</bundle>`);
  const spk = `<signedPreKeyPublic signedPreKeyId="1">${toBase64(key(2))}</signedPreKeyPublic>`;
  const sig = `<signedPreKeySignature>${toBase64(key(3, 64))}</signedPreKeySignature>`;
  const idk = `<identityKey>${toBase64(key(1, 33))}</identityKey>`;
  // Not a bundle at all.
  assert.equal(parseBundle(el(`<list xmlns="${NS}"><device id="1"/></list>`)), null);
  // Missing required parts.
  assert.equal(parseBundle(B(`${spk}${sig}`)), null); // no identityKey
  assert.equal(parseBundle(B(`${sig}${idk}`)), null); // no signedPreKeyPublic
  assert.equal(parseBundle(B(`${spk}${idk}`)), null); // no signature
  assert.equal(parseBundle(B("")), null);
  // A malformed signed prekey id.
  assert.equal(parseBundle(B(`<signedPreKeyPublic signedPreKeyId="x">AAA=</signedPreKeyPublic>${sig}${idk}`)), null);
  // Undecodable base64 in a required part.
  assert.equal(parseBundle(B(`<signedPreKeyPublic signedPreKeyId="1">!!!</signedPreKeyPublic>${sig}${idk}`)), null);
  assert.equal(parseBundle(B(`${spk}<signedPreKeySignature>!!!</signedPreKeySignature>${idk}`)), null);
  assert.equal(parseBundle(B(`${spk}${sig}<identityKey>!!!</identityKey>`)), null);
  // An empty required part is as malformed as a missing one.
  assert.equal(parseBundle(B(`${spk}${sig}<identityKey></identityKey>`)), null);
});

test("invalid prekey entries are skipped, the rest stay usable", () => {
  const bundle = el(
    `<bundle xmlns="${NS}">` +
      `<signedPreKeyPublic signedPreKeyId="1">${toBase64(key(2))}</signedPreKeyPublic>` +
      `<signedPreKeySignature>${toBase64(key(3, 64))}</signedPreKeySignature>` +
      `<identityKey>${toBase64(key(1, 33))}</identityKey>` +
      `<prekeys>` +
      `<preKeyPublic preKeyId="1">!!!</preKeyPublic>` +
      `<preKeyPublic preKeyId="x">${toBase64(key(4))}</preKeyPublic>` +
      `<preKeyPublic/>` +
      `<preKeyPublic preKeyId="9">${toBase64(key(5))}</preKeyPublic>` +
      `</prekeys>` +
      `</bundle>`,
  );
  const parsed = parseBundle(bundle)!;
  assert.deepEqual(
    parsed.preKeys.map((k) => k.keyId),
    [9],
  );
});

test("a device list round-trips, deduplicated, in the legacy wire format", () => {
  const built = buildDeviceList([42, 7, 42]);
  assert.equal(built.toString(), `<list xmlns="${NS}"><device id="42"/><device id="7"/></list>`);
  assert.deepEqual(parseDeviceList(built), [42, 7]);
});

test("a malformed or missing device list parses to empty, never throws", () => {
  assert.deepEqual(
    parseDeviceList(el(`<list xmlns="${NS}"><device id="1"/><device/><device id="x"/><device id="2"/><device id="1"/></list>`)),
    [1, 2],
  );
  // Wrong element or wrong namespace: empty — "no OMEMO devices", which
  // callers treat as not-OMEMO-capable.
  assert.deepEqual(parseDeviceList(el(`<bundle xmlns="${NS}"/>`)), []);
  assert.deepEqual(parseDeviceList(el(`<list xmlns="urn:xmpp:omemo:2"><device id="1"/></list>`)), []);
  assert.deepEqual(parseDeviceList(el(`<list xmlns="${NS}"/>`)), []);
});

test("an encrypted message round-trips through build→parse, in the legacy wire format", () => {
  const msg = {
    senderDeviceId: 111,
    keys: [
      { deviceId: 222, prekey: true, data: key(6, 32) },
      { deviceId: 333, prekey: false, data: key(7, 32) },
    ],
    iv: key(8, 12),
    payload: key(9, 40),
  };
  const built = buildEncryptedMessage(msg);
  // Lock the exact wire layout, not just a round-trip through our own
  // parser: `prekey` must be the literal attribute "true", present only
  // when true (never prekey="false"), the way deployed clients expect.
  assert.equal(
    built.toString(),
    `<encrypted xmlns="${NS}">` +
      `<header sid="111">` +
      `<key rid="222" prekey="true">${toBase64(key(6, 32))}</key>` +
      `<key rid="333">${toBase64(key(7, 32))}</key>` +
      `<iv>${toBase64(key(8, 12))}</iv>` +
      `</header>` +
      `<payload>${toBase64(key(9, 40))}</payload>` +
      `</encrypted>`,
  );
  const parsed = parseEncryptedMessage(built)!;
  assert.equal(parsed.senderDeviceId, 111);
  assert.deepEqual(
    parsed.keys.map((k) => [k.deviceId, k.prekey]),
    [
      [222, true],
      [333, false],
    ],
  );
  assert.deepEqual(new Uint8Array(parsed.keys[0].data), new Uint8Array(key(6, 32)));
  assert.deepEqual(new Uint8Array(parsed.iv), new Uint8Array(key(8, 12)));
  assert.deepEqual(new Uint8Array(parsed.payload), new Uint8Array(key(9, 40)));
});

test("an encrypted message with no keys for us still parses (empty keys array)", () => {
  const built = buildEncryptedMessage({ senderDeviceId: 1, keys: [], iv: key(1, 12), payload: key(2, 10) });
  const parsed = parseEncryptedMessage(built)!;
  assert.deepEqual(parsed.keys, []);
});

test("a malformed or missing encrypted message parses to null, never throws", () => {
  const E = (inner: string) => el(`<encrypted xmlns="${NS}">${inner}</encrypted>`);
  const iv = `<iv>${toBase64(key(1, 12))}</iv>`;
  const payload = `<payload>${toBase64(key(2, 10))}</payload>`;
  // Not an <encrypted> at all.
  assert.equal(parseEncryptedMessage(el(`<bundle xmlns="${NS}"/>`)), null);
  // Missing header, missing sid, missing iv, missing payload.
  assert.equal(parseEncryptedMessage(E(payload)), null);
  assert.equal(parseEncryptedMessage(E(`<header>${iv}</header>${payload}`)), null);
  assert.equal(parseEncryptedMessage(E(`<header sid="1"></header>${payload}`)), null);
  assert.equal(parseEncryptedMessage(E(`<header sid="1">${iv}</header>`)), null);
  // A malformed sid, and undecodable base64 in iv/payload.
  assert.equal(parseEncryptedMessage(E(`<header sid="x">${iv}</header>${payload}`)), null);
  assert.equal(parseEncryptedMessage(E(`<header sid="1"><iv>!!!</iv></header>${payload}`)), null);
  assert.equal(parseEncryptedMessage(E(`<header sid="1">${iv}</header><payload>!!!</payload>`)), null);
});

test("an encrypted message with one malformed key entry skips it, keeps the rest", () => {
  const built = el(
    `<encrypted xmlns="${NS}">` +
      `<header sid="1">` +
      `<key rid="2">${toBase64(key(1, 32))}</key>` +
      `<key>${toBase64(key(2, 32))}</key>` + // missing rid
      `<key rid="3"></key>` + // empty text
      `<iv>${toBase64(key(3, 12))}</iv>` +
      `</header>` +
      `<payload>${toBase64(key(4, 10))}</payload>` +
      `</encrypted>`,
  );
  const parsed = parseEncryptedMessage(built)!;
  assert.deepEqual(parsed.keys.map((k) => k.deviceId), [2]);
});