import { describe, expect, it } from "vitest";
import { generateKeyPair } from "@stablelib/x25519";
import { noiseXXInitiator, noiseXXResponder } from "./noise-xx";

function memoryPair() {
  const a: Uint8Array[] = [];
  const b: Uint8Array[] = [];
  const aR: Array<(v: Uint8Array) => void> = [];
  const bR: Array<(v: Uint8Array) => void> = [];
  const tA = {
    async send(m: Uint8Array) {
      const r = bR.shift();
      if (r) r(m); else b.push(m);
    },
    receive(): Promise<Uint8Array> {
      const v = a.shift();
      if (v) return Promise.resolve(v);
      return new Promise<Uint8Array>((resolve) => aR.push(resolve));
    },
  };
  const tB = {
    async send(m: Uint8Array) {
      const r = aR.shift();
      if (r) r(m); else a.push(m);
    },
    receive(): Promise<Uint8Array> {
      const v = b.shift();
      if (v) return Promise.resolve(v);
      return new Promise<Uint8Array>((resolve) => bR.push(resolve));
    },
  };
  return [tA, tB] as const;
}

describe("Noise XX mutual authentication", () => {
  it("completes the three-message handshake and exchanges static keys", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const [tA, tB] = memoryPair();

    const aPromise = noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive);
    const bPromise = noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive);

    const [a, b] = await Promise.all([aPromise, bPromise]);
    expect(a.remoteStaticPublicKey).toEqual(bob.publicKey);
    expect(b.remoteStaticPublicKey).toEqual(alice.publicKey);
    expect(a.handshakeHash).toEqual(b.handshakeHash);
  });

  it("transport frames decrypt in both directions and are bound to the handshake", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const [tA, tB] = memoryPair();

    const [a, b] = await Promise.all([
      noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive),
      noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive),
    ]);

    const ct1 = a.encrypt(new TextEncoder().encode("hello bob"));
    expect(new TextDecoder().decode(b.decrypt(ct1))).toBe("hello bob");

    const ct2 = b.encrypt(new TextEncoder().encode("hello alice"));
    expect(new TextDecoder().decode(a.decrypt(ct2))).toBe("hello alice");

    // Tampering with one byte breaks the MAC.
    const ct3 = a.encrypt(new TextEncoder().encode("not delivered"));
    ct3[0] ^= 1;
    expect(() => b.decrypt(ct3)).toThrow();
  });

  it("an attacker who swaps the responder static cannot complete msg2 auth", async () => {
    const alice = generateKeyPair();
    const bobReal = generateKeyPair();
    const bobImpostor = generateKeyPair();
    const [tA, tB] = memoryPair();

    // Responder uses impostor private key but advertises real pubkey to initiator.
    // Real XX has no out-of-band pubkey check; this test verifies the registry
    // check that the caller MUST perform on `remoteStaticPublicKey`.
    const aPromise = noiseXXInitiator(alice.secretKey, alice.publicKey, tA.send, tA.receive);
    const bPromise = noiseXXResponder(bobImpostor.secretKey, bobImpostor.publicKey, tB.send, tB.receive);
    const [a] = await Promise.all([aPromise, bPromise]);
    // The "swap" is the caller seeing a remote static they did not expect.
    expect(a.remoteStaticPublicKey).not.toEqual(bobReal.publicKey);
    expect(a.remoteStaticPublicKey).toEqual(bobImpostor.publicKey);
  });

  // PCH-01-003: the pre-msg3 verifier must abort BEFORE the initiator transmits
  // its own static (msg3), so a probed-but-unexpected LAN host never learns our
  // long-term sync identity.
  it("initiator_aborts_before_msg3_on_unexpected_responder", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const expected = generateKeyPair(); // the peer we WANTED — not bob
    const [tA, tB] = memoryPair();

    let sendCount = 0;
    const countingSend = async (m: Uint8Array) => {
      sendCount++;
      await tA.send(m);
    };

    const aPromise = noiseXXInitiator(
      alice.secretKey,
      alice.publicKey,
      countingSend,
      tA.receive,
      (remoteStatic) => {
        const matches =
          remoteStatic.length === expected.publicKey.length &&
          remoteStatic.every((v, i) => v === expected.publicKey[i]);
        if (!matches) throw new Error("paired-device static key mismatch");
      },
    );
    // Honest responder advertising bob's static; it will hang waiting for msg3
    // that never arrives, so we don't await it.
    const bPromise = noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive);
    bPromise.catch(() => {});

    await expect(aPromise).rejects.toThrow("paired-device static key mismatch");
    // Load-bearing assertion: only msg1 was sent. msg3 (our static) never left.
    expect(sendCount).toBe(1);
  });

  it("matching_pin_completes_handshake_and_sends_msg3", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const [tA, tB] = memoryPair();

    let sendCount = 0;
    const countingSend = async (m: Uint8Array) => {
      sendCount++;
      await tA.send(m);
    };

    const [a] = await Promise.all([
      noiseXXInitiator(alice.secretKey, alice.publicKey, countingSend, tA.receive, (remoteStatic) => {
        const matches = remoteStatic.every((v, i) => v === bob.publicKey[i]);
        if (!matches) throw new Error("paired-device static key mismatch");
      }),
      noiseXXResponder(bob.secretKey, bob.publicKey, tB.send, tB.receive),
    ]);

    expect(a.remoteStaticPublicKey).toEqual(bob.publicKey);
    // msg1 + msg3 both sent on the honest path — handshake completed.
    expect(sendCount).toBe(2);
  });
});