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("E6: every reconnect derives fresh keys, so one session never exposes another", async () => {
    // Peach sync sessions are one-shot: ~10 frames, a 30s coordinator timeout,
    // then close. There is no long-lived session to rekey inside, and each
    // reconnect runs a full KK handshake with fresh ephemerals — which gives
    // strictly stronger post-compromise security than intra-session Rekey()
    // would. This pins that property so a future change cannot quietly reuse
    // ephemerals and turn every past session's traffic into one key's problem.
    const initiator = generateKeyPair();
    const responder = generateKeyPair();

    async function handshake() {
      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,
        ),
      ]);
      return { a, b };
    }

    const first = await handshake();
    const second = await handshake();

    // Same long-term statics, different session — the transcripts must differ.
    expect(bytesToHex(first.a.handshakeHash)).not.toBe(
      bytesToHex(second.a.handshakeHash),
    );

    // And the keys must genuinely differ: a frame from session two must not be
    // readable by session one.
    const frame = second.a.encrypt(new TextEncoder().encode("session two"));
    expect(() => first.b.decrypt(frame)).toThrow();
  });

  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(
      "ac3dd2e4c40290fab2a0192b0b33b37bb8c31396ffec23271ffa2a99ef3a0f75",
    );
    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(
      /decrypt/i,
    );
  });
});