// Run: npm test   (Node strips the TypeScript types itself)
// Registers globalThis.indexedDB before anything imports src/store/idb.ts —
// Node has no IndexedDB of its own.
import "fake-indexeddb/auto";
import { test } from "node:test";
import assert from "node:assert/strict";
import { xml, type Element } from "@xmpp/client";
import { OmemoDB, type StoredPreKey } from "../src/crypto/omemo-db.ts";
import { OmemoStore } from "../src/crypto/omemo-store.ts";
import {
  bootstrapIdentity,
  decrypt,
  encryptFor,
  getTrustStatus,
  keyFingerprintHex,
  markVerified,
  omemoSupported,
  retractOwnDevices,
} from "../src/crypto/omemo.ts";
import { buildBundle, buildDeviceList, buildEncryptedMessage } from "../src/crypto/omemo-wire.ts";
import { composeReply, replyBody } from "../src/xmpp/replies.ts";

// A distinct account per test: fake-indexeddb keeps its databases per process.
const accounts = ["lukas@0ff.space", "anna@0ff.space", "bob@0ff.space", "clara@0ff.space", "dora@0ff.space", "ewa@0ff.space"];
const [lukas, anna, bob, clara, dora, ewa] = accounts;
// Distinct non-zero identity keys, for the trust-change checks.
const identityKey = (byte: number) => {
  const b = new ArrayBuffer(33);
  new Uint8Array(b)[0] = byte;
  return b;
};

test("identity bootstrap is idempotent across repeated opens (relogin)", async () => {
  const first = await bootstrapIdentity(lukas);
  // A second open of omemo-db simulates logging back in on the same browser:
  // the stored identity must come back unchanged, never freshly generated.
  const second = await bootstrapIdentity(lukas);
  assert.equal(second.registrationId, first.registrationId);
  assert.equal(second.deviceId, first.deviceId);
  assert.deepEqual(second.keyPair, first.keyPair);
  // Third call, straight through a fresh OmemoDB handle, same identity again.
  const db = await OmemoDB.open(lukas);
  try {
    assert.deepEqual(await db.getIdentity(), first);
  } finally {
    db.close();
  }
});

test("keys round-trip through IndexedDB as plain structured-cloneable values", async () => {
  await bootstrapIdentity(anna);
  const db = await OmemoDB.open(anna);
  try {
    // Class instances and functions throw in structuredClone; the stored rows
    // must be plain cloneable values, or later logout/relogin reads break.
    const identity = await db.getIdentity();
    assert.ok(identity);
    const signedPreKeys = await db.getSignedPreKeys();
    assert.equal(signedPreKeys.length, 1);
    const preKeys = await db.getPreKeys();
    assert.equal(preKeys.length, 100);
    assert.ok(preKeys.every((k) => !k.used));
    for (const row of [identity, signedPreKeys[0], preKeys[0], preKeys[99]]) {
      assert.deepEqual(structuredClone(row), row);
    }
  } finally {
    db.close();
  }
});

test("OmemoStore adapts prekeys, signed prekeys, sessions and trust to OmemoDB", async () => {
  const identity = await bootstrapIdentity(bob);
  const db = await OmemoDB.open(bob);
  try {
    const store = new OmemoStore(db, "eu.siacs.conversations.axolotl");
    assert.deepEqual(await store.getIdentityKeyPair(), identity.keyPair);
    assert.equal(await store.getLocalRegistrationId(), identity.registrationId);

    // One-time prekeys: store, load, and removed-after-use never reload.
    // (Ids outside bootstrap's 1..100, so these rows are the test's own.)
    const keyPair = { pubKey: identityKey(1), privKey: identityKey(2) };
    await store.storePreKey(500, keyPair);
    assert.deepEqual((await store.loadPreKey(500))?.keyPair, keyPair);
    assert.equal(await store.loadPreKey(501), undefined);
    await store.removePreKey(500);
    assert.equal(await store.loadPreKey(500), undefined);

    // Signed prekeys: the bootstrap ones load through the adapter.
    const spk = await store.loadSignedPreKey(1);
    assert.ok(spk);
    assert.deepEqual(spk.keyPair, (await db.getSignedPreKey(1))?.keyPair);

    // Sessions: plain serialized strings, keyed by the library's own
    // `${bareJid}.${deviceId}` address, dotted domains included.
    const record = JSON.stringify({ opaque: "ratchet state" });
    await store.storeSession("anna@0ff.space.co.uk.42", record);
    assert.equal(await store.loadSession("anna@0ff.space.co.uk.42"), record);
    assert.equal(await store.loadSession("anna@0ff.space.co.uk.43"), undefined);
    await store.storeSession("anna@0ff.space.co.uk.43", record);
    await store.removeAllSessions("anna@0ff.space.co.uk");
    assert.equal(await store.loadSession("anna@0ff.space.co.uk.42"), undefined);
    assert.equal(await store.loadSession("anna@0ff.space.co.uk.43"), undefined);

    // Trust: silent TOFU — first key pins silently; isTrustedIdentity never
    // blocks, even against a changed key (OMEMO-PLAN.md, Trust model: warn,
    // don't block — the library turns `false` into a thrown "Identity key
    // changed" otherwise). saveIdentity is what actually detects and
    // reports a change.
    const first = identityKey(5);
    assert.equal(await store.isTrustedIdentity("anna@0ff.space.42", first, 2), true);
    assert.equal(await store.saveIdentity("anna@0ff.space.42", first), false);
    assert.equal(await store.isTrustedIdentity("anna@0ff.space.42", first, 2), true);
    const changed = identityKey(6);
    assert.equal(await store.isTrustedIdentity("anna@0ff.space.42", changed, 2), true);
    assert.equal(await store.saveIdentity("anna@0ff.space.42", changed), true);
    assert.deepEqual(await store.loadIdentityKey("anna@0ff.space.42"), changed);
  } finally {
    db.close();
  }
});

test("one store per namespace: same address, never the same session or trust", async () => {
  await bootstrapIdentity(clara);
  const db = await OmemoDB.open(clara);
  try {
    const legacy = new OmemoStore(db, "eu.siacs.conversations.axolotl");
    const omemo2 = new OmemoStore(db, "urn:xmpp:omemo:2");
    const address = "anna@0ff.space.42";
    await legacy.storeSession(address, "legacy ratchet");
    await omemo2.storeSession(address, "omemo2 ratchet");
    assert.equal(await legacy.loadSession(address), "legacy ratchet");
    assert.equal(await omemo2.loadSession(address), "omemo2 ratchet");
    await legacy.removeAllSessions("anna@0ff.space");
    assert.equal(await legacy.loadSession(address), undefined);
    assert.equal(await omemo2.loadSession(address), "omemo2 ratchet");
    // Pinned trust is per namespace too: legacy's pin is untouched by
    // omemo2's calls, and omemo2 never pinned anything at all. (Both
    // isTrustedIdentity calls return true regardless — it never blocks.)
    const first = identityKey(5);
    await legacy.saveIdentity(address, first);
    assert.equal(await omemo2.isTrustedIdentity(address, identityKey(9), 2), true);
    assert.equal(await legacy.isTrustedIdentity(address, identityKey(9), 2), true);
    assert.deepEqual(await legacy.loadIdentityKey(address), first);
    assert.equal(await omemo2.loadIdentityKey(address), undefined);
  } finally {
    db.close();
  }
});

test("trust status: tofu -> changed -> verified, and a further change re-flags it", async () => {
  await bootstrapIdentity(ewa + "-trust");
  const db = await OmemoDB.open(ewa + "-trust");
  try {
    const store = new OmemoStore(db, "eu.siacs.conversations.axolotl");
    const peer = "bob@0ff.space";
    const deviceId = 7;
    const address = `${peer}.${deviceId}`;

    assert.equal(await getTrustStatus(db, peer, deviceId), undefined); // nothing pinned yet
    await store.saveIdentity(address, identityKey(1));
    assert.equal((await getTrustStatus(db, peer, deviceId))?.status, "tofu");

    await markVerified(db, peer, deviceId);
    assert.equal((await getTrustStatus(db, peer, deviceId))?.status, "verified");

    // A real key change always re-flags "changed", even over "verified" —
    // verifying key A says nothing about a later, different key B.
    await store.saveIdentity(address, identityKey(2));
    assert.equal((await getTrustStatus(db, peer, deviceId))?.status, "changed");
    // markVerified on a device with no pinned key at all is a safe no-op.
    await markVerified(db, "nobody@0ff.space", 1);
    assert.equal(await getTrustStatus(db, "nobody@0ff.space", 1), undefined);
  } finally {
    db.close();
  }
});

test("keyFingerprintHex is a stable, readable, grouped hex string", () => {
  const key = identityKey(10);
  const hex = keyFingerprintHex(key);
  assert.equal(hex, keyFingerprintHex(key)); // stable for the same input
  assert.match(hex, /^[0-9A-F]{1,8}( [0-9A-F]{1,8})*$/); // uppercase, grouped in 8s
  assert.notEqual(keyFingerprintHex(identityKey(11)), hex); // distinct keys, distinct output
});

test("keyFingerprintHex drops the 0x05 type byte, as Monal and Conversations do", () => {
  // A real Curve25519 identity key: the 0x05 DJB type byte then 32 key
  // bytes. Confirmed 2026-10-07 against Monal on iPhone, which displayed a
  // fingerprint starting DA02CF0F for a stored key starting 05DA02CF —
  // i.e. the type byte is not shown. Rendering all 33 bytes put every
  // character one byte out of step with the other app, so a user comparing
  // them out of band would read "doesn't match" for identical keys.
  const bytes = new Uint8Array(33);
  bytes[0] = 0x05;
  for (let i = 1; i < 33; i++) bytes[i] = i;
  const hex = keyFingerprintHex(bytes.buffer);
  assert.equal(hex, "01020304 05060708 090A0B0C 0D0E0F10 11121314 15161718 191A1B1C 1D1E1F20");
  // Exactly 8 groups of 8: no stray short group, which is what a trailing
  // 33rd byte produced and the clearest tell that the format was wrong.
  assert.equal(hex.split(" ").length, 8);
  assert.ok(hex.split(" ").every((g) => g.length === 8));
  assert.equal(hex.replace(/ /g, "").length, 64);

  // A key that is not 33 bytes, or does not start with 0x05, is rendered
  // whole: never silently eat a byte of something we don't recognise.
  const odd = new Uint8Array([0x04, 0xaa, 0xbb]);
  assert.equal(keyFingerprintHex(odd.buffer), "04AABB");
});

test("omemoSupported reports true when WebAssembly actually works (this environment)", async () => {
  // Node has no CSP, so the real, intended check — does WebAssembly.instantiate
  // throw? — must pass here. The *blocked* branch (a CSP throwing) isn't
  // separately unit-tested: omemoSupported() memoizes its result at module
  // level for the lifetime of the page (deliberately — WASM support doesn't
  // change mid-session), which in this test file's single process would
  // permanently poison every later test that needs real OMEMO crypto
  // (encryptFor/decrypt) if exercised here first. The gating logic itself —
  // `if (!(await omemoSupported())) return [];` in cachedDeviceList — is
  // trivial enough to trust from code review; this test only confirms the
  // non-degraded path isn't itself broken.
  assert.equal(await omemoSupported(), true);
});

// A fake PEP service for `encryptFor`/`decrypt`'s own fetch calls
// (fetchDeviceList/fetchBundle), so Phase 3's send-path crypto can be
// round-tripped end to end without a real XMPP connection — the only way
// to verify `encryptFor`'s wire format locally, short of a live Monal.
const NS_PUBSUB = "http://jabber.org/protocol/pubsub";

class FakeOmemoNetwork {
  #devicelists = new Map<string, number[]>();
  #bundles = new Map<string, Element>(); // key: `${bareJid}:${deviceId}`

  /** Publish one device's bundle directly — e.g. a second device on an
   * already-registered account (its own OmemoDB, same bare JID). */
  registerDevice(bareJid: string, deviceId: number, bundle: Element): void {
    this.#bundles.set(`${bareJid}:${deviceId}`, bundle);
    this.#devicelists.set(bareJid, [...(this.#devicelists.get(bareJid) ?? []), deviceId]);
  }

  /** Bootstrap a fresh identity for `bareJid` in `db` and publish its bundle. */
  async register(db: OmemoDB, bareJid: string): Promise<number> {
    const identity = await bootstrapIdentity(bareJid);
    const [signedPreKey] = await db.getSignedPreKeys();
    const preKeys = (await db.getPreKeys()).filter((k) => !k.used);
    const bundle = buildBundle(
      identity.keyPair.pubKey,
      { keyId: signedPreKey.keyId, publicKey: signedPreKey.keyPair.pubKey, signature: signedPreKey.signature! },
      preKeys.map((k: StoredPreKey) => ({ keyId: k.keyId, publicKey: k.keyPair.pubKey })),
    );
    this.registerDevice(bareJid, identity.deviceId, bundle);
    return identity.deviceId;
  }

  iqCaller = {
    get: async (stanza: Element, to: string): Promise<Element> => {
      const node = stanza.getChild("items")?.attrs.node as string;
      let payload: Element | undefined;
      if (node === "eu.siacs.conversations.axolotl.devicelist") {
        payload = buildDeviceList(this.#devicelists.get(to) ?? []);
      } else {
        const m = /^eu\.siacs\.conversations\.axolotl\.bundles:(\d+)$/.exec(node ?? "");
        if (m) payload = this.#bundles.get(`${to}:${m[1]}`);
      }
      if (!payload) {
        const err = new Error("item-not-found") as Error & { condition: string };
        err.condition = "item-not-found";
        throw err;
      }
      return xml("pubsub", { xmlns: NS_PUBSUB }, xml("items", { node }, xml("item", { id: "current" }, payload)));
    },
    // Only what retractOwnDevices needs: republishing a device list, and
    // deleting a bundle node (pubsub#owner).
    set: async (stanza: Element, to: string): Promise<Element> => {
      const publish = stanza.getChild("publish");
      if (publish?.attrs.node === "eu.siacs.conversations.axolotl.devicelist") {
        const list = publish.getChild("item")?.getChild("list");
        const ids = (list?.getChildren("device") ?? []).map((d) => Number(d.attrs.id));
        this.#devicelists.set(to, ids);
        this.publishedItemIds.push(String(publish.getChild("item")?.attrs.id));
      }
      const del = stanza.getChild("delete");
      if (del?.attrs.node) this.deletedNodes.push(String(del.attrs.node));
      return xml("iq");
    },
  };

  /** Item ids the device list was republished under, in order. */
  publishedItemIds: string[] = [];
  /** Bundle nodes retractOwnDevices asked the server to delete. */
  deletedNodes: string[] = [];

  devicesOf(bareJid: string): number[] {
    return this.#devicelists.get(bareJid) ?? [];
  }
}

test("encryptFor/decrypt round-trip real plaintext through the Double Ratchet + AES-GCM, no network", async () => {
  const net = new FakeOmemoNetwork();
  const doraDb = await OmemoDB.open(dora);
  const ewaDb = await OmemoDB.open(ewa);
  try {
    const doraDeviceId = await net.register(doraDb, dora);
    const ewaDeviceId = await net.register(ewaDb, ewa);

    const envelope = await encryptFor(net as any, doraDb, dora, doraDeviceId, ewa, "hello ewa, this is dora");
    // The wire format never carries the plaintext or the raw AES key.
    assert.equal(envelope.keys.length, 1); // ewa's one device only; dora has no other device
    assert.equal(envelope.keys[0].deviceId, ewaDeviceId);
    assert.equal(envelope.keys[0].prekey, true); // first message on a brand-new session

    const plaintext = await decrypt(ewaDb, dora, doraDeviceId, ewaDeviceId, envelope);
    assert.equal(new TextDecoder().decode(plaintext), "hello ewa, this is dora");
    const quote = composeReply("encrypted answer 🎉", { key: "quote", peer: ewa, id: "original-id", to: ewa, text: "private original 🐱", encrypted: true });
    const replyEnvelope = await encryptFor(net as any, doraDb, dora, doraDeviceId, ewa, quote.body);
    const stanza = xml("message", null, xml("body", null, "This message is encrypted"), buildEncryptedMessage(replyEnvelope), ...quote.elements);
    assert.ok(!stanza.toString().includes("private original"));
    assert.ok(!stanza.toString().includes("encrypted answer"));
    const monalText = new TextDecoder().decode(await decrypt(ewaDb, dora, doraDeviceId, ewaDeviceId, replyEnvelope));
    assert.equal(monalText, quote.body); // Monal's legacy parser displays this text unchanged.
    assert.deepEqual(replyBody(stanza, monalText, { id: "original-id", to: ewa }), { body: "encrypted answer 🎉", reply: quote.reply });
  } finally {
    doraDb.close();
    ewaDb.close();
  }
});

test("encryptFor wraps a key for the sender's own other device too", async () => {
  const net = new FakeOmemoNetwork();
  const account = dora + "2";
  const doraDb = await OmemoDB.open(account); // the "web" device sending the message
  const ewaDb = await OmemoDB.open(ewa + "2");
  const doraPhoneDb = await OmemoDB.open(account + "-phone"); // a second device, same account
  try {
    const doraDeviceId = await net.register(doraDb, account);
    const ewaDeviceId = await net.register(ewaDb, ewa + "2");

    // The phone's identity lives in its own OmemoDB (a distinct device), but
    // is registered under dora's *same bare JID* in the fake network — what
    // a real second device publishing its own bundle to the shared
    // devicelist node looks like.
    const phoneIdentity = await bootstrapIdentity(account + "-phone");
    const [phoneSignedPreKey] = await doraPhoneDb.getSignedPreKeys();
    const phonePreKeys = (await doraPhoneDb.getPreKeys()).filter((k) => !k.used);
    const phoneBundle = buildBundle(
      phoneIdentity.keyPair.pubKey,
      { keyId: phoneSignedPreKey.keyId, publicKey: phoneSignedPreKey.keyPair.pubKey, signature: phoneSignedPreKey.signature! },
      phonePreKeys.map((k) => ({ keyId: k.keyId, publicKey: k.keyPair.pubKey })),
    );
    net.registerDevice(account, phoneIdentity.deviceId, phoneBundle);

    const envelope = await encryptFor(net as any, doraDb, account, doraDeviceId, ewa + "2", "hi from the web client");
    // Two recipients: ewa's device and dora's own phone — never dora's own
    // sending device.
    assert.deepEqual(
      envelope.keys.map((k) => k.deviceId).sort((a, b) => a - b),
      [ewaDeviceId, phoneIdentity.deviceId].sort((a, b) => a - b),
    );

    // The phone decrypts its own wrapped key and reads back the exact text.
    const plaintext = await decrypt(doraPhoneDb, account, doraDeviceId, phoneIdentity.deviceId, envelope);
    assert.equal(new TextDecoder().decode(plaintext), "hi from the web client");
  } finally {
    doraDb.close();
    ewaDb.close();
    doraPhoneDb.close();
  }
});

test("retractOwnDevices drops dead devices but never the one in use", async () => {
  const net = new FakeOmemoNetwork();
  const me = "gil@example.invalid";
  // Three stale ids from earlier browser profiles, plus the device in use.
  const inUse = 4242;
  net.registerDevice(me, 111, xml("bundle"));
  net.registerDevice(me, 222, xml("bundle"));
  net.registerDevice(me, inUse, xml("bundle"));

  // `inUse` is passed in the removal list too: a caller must not be able to
  // retract the device it is running on, which would go silently unreachable.
  const removed = await retractOwnDevices(net as any, me, [111, 222, inUse], inUse);

  assert.deepEqual(removed.sort(), [111, 222]);
  assert.deepEqual(net.devicesOf(me), [inUse]);
  // Republished in place, under the item id every deployed client uses.
  assert.deepEqual(net.publishedItemIds, ["current"]);
  // Each retracted device's bundle node is cleaned up as well.
  assert.deepEqual(net.deletedNodes.sort(), [
    "eu.siacs.conversations.axolotl.bundles:111",
    "eu.siacs.conversations.axolotl.bundles:222",
  ]);
});

test("retractOwnDevices is a no-op when there is nothing stale to remove", async () => {
  const net = new FakeOmemoNetwork();
  const me = "hana@example.invalid";
  net.registerDevice(me, 777, xml("bundle"));

  assert.deepEqual(await retractOwnDevices(net as any, me, [777], 777), []);
  assert.deepEqual(await retractOwnDevices(net as any, me, [], 777), []);
  // Nothing published, nothing deleted: an unchanged list must not be
  // republished, which would churn every contact's devicelist subscription.
  assert.deepEqual(net.publishedItemIds, []);
  assert.deepEqual(net.deletedNodes, []);
  assert.deepEqual(net.devicesOf(me), [777]);
});

test("device additions and removals change wrapped targets while preserving decryptable ratchets", async () => {
  const net = new FakeOmemoNetwork();
  const account = "target-a@0ff.space", peer = "target-b@0ff.space";
  const a = await OmemoDB.open(account), b = await OmemoDB.open(peer);
  const bNew = await OmemoDB.open(peer + "-new"), aNew = await OmemoDB.open(account + "-new");
  const { acceptDeviceListEvent, cachedDeviceList, invalidateDeviceLists } = await import('../src/crypto/omemo.ts');
  const add = async (db: OmemoDB, storage: string, publisher: string) => {
    const identity = await bootstrapIdentity(storage);
    const spk = (await db.getSignedPreKeys())[0];
    net.registerDevice(publisher, identity.deviceId, buildBundle(identity.keyPair.pubKey,
      {keyId:spk.keyId, publicKey:spk.keyPair.pubKey, signature:spk.signature!},
      (await db.getPreKeys()).map(k => ({keyId:k.keyId, publicKey:k.keyPair.pubKey}))));
    return identity.deviceId;
  };
  try {
    const aid = await net.register(a, account), bid = await net.register(b, peer);
    await cachedDeviceList(net as any, a, peer);
    const newBid = await add(bNew, peer + '-new', peer);
    await acceptDeviceListEvent(net as any, a, peer, [bid,newBid], true);
    const first = await encryptFor(net as any, a, account, aid, peer, 'new B can read');
    assert.equal(new TextDecoder().decode(await decrypt(bNew, account, aid, newBid, first)), 'new B can read');
    const newAid = await add(aNew, account + '-new', account);
    await acceptDeviceListEvent(net as any, a, account, [aid,newAid], true);
    await acceptDeviceListEvent(net as any, a, peer, [newBid], true);
    const second = await encryptFor(net as any, a, account, aid, peer, 'removed B omitted');
    assert.deepEqual(second.keys.map(k=>k.deviceId).sort(), [newBid,newAid].sort());
    assert.equal(new TextDecoder().decode(await decrypt(bNew, account, aid, newBid, second)), 'removed B omitted');
    assert.equal(new TextDecoder().decode(await decrypt(aNew, account, aid, newAid, second)), 'removed B omitted');
    invalidateDeviceLists(net as any);
    const third = await encryptFor(net as any, a, account, aid, peer, 'reconnected');
    assert.equal(new TextDecoder().decode(await decrypt(bNew, account, aid, newBid, third)), 'reconnected');
  } finally {a.close();b.close();bNew.close();aNew.close();}
});