import { describe, expect, it } from "vitest";
import { generateKeyPair, generateKeyPairFromSeed } from "@stablelib/x25519";
import {
  bytesToHex,
  noiseKKInitiator,
  noiseKKResponder,
} from "./noise-js";

function channels() {
  const initiatorToResponder: Uint8Array[] = [];
  const responderToInitiator: Uint8Array[] = [];
  const drain = async (queue: Uint8Array[]) => {
    while (queue.length === 0) await new Promise<void>((resolve) => setTimeout(resolve, 0));
    return queue.shift()!;
  };
  return {
    initiatorSend: async (frame: Uint8Array) => { initiatorToResponder.push(frame); },
    initiatorReceive: () => drain(responderToInitiator),
    responderSend: async (frame: Uint8Array) => { responderToInitiator.push(frame); },
    responderReceive: () => drain(initiatorToResponder),
  };
}

describe("Noise KK reconnect", () => {
  it("mutually authenticates known static keys and carries transport frames", async () => {
    const initiator = generateKeyPair();
    const responder = generateKeyPair();
    const wire = channels();
    const [a, b] = await Promise.all([
      noiseKKInitiator(
        initiator.secretKey,
        initiator.publicKey,
        responder.publicKey,
        wire.initiatorSend,
        wire.initiatorReceive,
      ),
      noiseKKResponder(
        responder.secretKey,
        responder.publicKey,
        initiator.publicKey,
        wire.responderSend,
        wire.responderReceive,
      ),
    ]);

    expect(bytesToHex(a.handshakeHash)).toBe(bytesToHex(b.handshakeHash));
    expect(new TextDecoder().decode(b.decrypt(a.encrypt(new TextEncoder().encode("one"))))).toBe("one");
    expect(new TextDecoder().decode(a.decrypt(b.encrypt(new TextEncoder().encode("two"))))).toBe("two");
  });

  it("pins the cross-platform deterministic handshake hash", async () => {
    const initiator = generateKeyPairFromSeed(new Uint8Array(32).fill(0x01));
    const responder = generateKeyPairFromSeed(new Uint8Array(32).fill(0x02));
    const initiatorEphemeral = generateKeyPairFromSeed(new Uint8Array(32).fill(0x03));
    const responderEphemeral = generateKeyPairFromSeed(new Uint8Array(32).fill(0x04));
    const wire = channels();
    const [a, b] = await Promise.all([
      noiseKKInitiator(
        initiator.secretKey,
        initiator.publicKey,
        responder.publicKey,
        wire.initiatorSend,
        wire.initiatorReceive,
        initiatorEphemeral,
      ),
      noiseKKResponder(
        responder.secretKey,
        responder.publicKey,
        initiator.publicKey,
        wire.responderSend,
        wire.responderReceive,
        responderEphemeral,
      ),
    ]);

    expect(bytesToHex(a.handshakeHash)).toBe(
      "e92392c530684df3eed06a76bd18da5c7a1887e88f5fdee758f9870dadd28182",
    );
    expect(bytesToHex(b.handshakeHash)).toBe(bytesToHex(a.handshakeHash));
  });

  it("hard-fails a mismatched expected initiator key without downgrade", async () => {
    const initiator = generateKeyPair();
    const attacker = generateKeyPair();
    const responder = generateKeyPair();
    const wire = channels();

    const initiatorHandshake = noiseKKInitiator(
      initiator.secretKey,
      initiator.publicKey,
      responder.publicKey,
      wire.initiatorSend,
      wire.initiatorReceive,
    );
    const responderHandshake = noiseKKResponder(
      responder.secretKey,
      responder.publicKey,
      attacker.publicKey,
      wire.responderSend,
      wire.responderReceive,
    );

    await expect(Promise.all([initiatorHandshake, responderHandshake])).rejects.toThrow(
      /Noise KK: msg1 auth failed/,
    );
  });
});