/**
 * Noise NK auth proof regression test.
 *
 * Pins a known input → output for the HMAC-SHA256 proof so any change to the
 * proof format is caught immediately. The pinned vector was captured against
 * the originally-shipped implementation; if the format ever needs to change,
 * regenerate both this vector and any peer that pairs against this build.
 */

import { describe, expect, it } from "vitest";
import { generateKeyPair as x25519GenerateKeyPair } from "@stablelib/x25519";

import {
  bytesToHex,
  computeAuthProof,
  noiseNKInitiator,
  noiseNKResponder,
  performInitiatorAuth,
  performResponderAuth,
} from "./noise-js";

// Pinned cross-check vector for HMAC-SHA256 auth proof.
const INITIATOR_PRIVATE = new Uint8Array(32).fill(0x01);
const NONCE = new Uint8Array(32).fill(0x03);
// PCH-01-003b: the proof now binds the Noise handshake hash. Pin a fixed one.
const HANDSHAKE_HASH = new Uint8Array(32).fill(0x04);

const EXPECTED_INITIATOR_PUBLIC =
  "a4e09292b651c278b9772c569f5fa9bb13d906b46ab68c9df9dc2b4409f8a209";
const EXPECTED_RESPONDER_PUBLIC =
  "ce8d3ad1ccb633ec7b70c17814a5c76ecd029685050d344745ba05870e587d59";
// Recaptured after PCH-01-003b prepended HANDSHAKE_HASH into the proof transcript.
const EXPECTED_PROOF =
  "d947d11a82892a65fdde94da5bbe6ca6034cf1f7e3268fc3ce5851b55f5813ce";

function derivePublic(_privateKey: Uint8Array): Uint8Array {
  const pair = x25519GenerateKeyPair();
  // We can't derive a public from an arbitrary private with @stablelib's
  // generateKeyPair (it samples a fresh secret). For the cross-check, the
  // expected public keys are pinned above so we just use them directly.
  void pair;
  throw new Error("derivePublic should not be called — use pinned constants");
}

describe("noise-js cross-check", () => {
  it("computes the pinned HMAC-SHA256 proof for the known input vector", () => {
    // Use the pinned public-key constants — they correspond to the X25519
    // public derived from the all-0x01 / all-0x02 private keys. Rederiving
    // them here would require a "derivePublic from raw scalar" helper that
    // @stablelib does not expose, so we trust the pinned values.
    void derivePublic;

    const initiatorPublic = hexToBytes(EXPECTED_INITIATOR_PUBLIC);
    const responderPublic = hexToBytes(EXPECTED_RESPONDER_PUBLIC);

    const proof = computeAuthProof(
      INITIATOR_PRIVATE,
      initiatorPublic,
      responderPublic,
      NONCE,
      HANDSHAKE_HASH,
    );

    expect(bytesToHex(proof)).toBe(EXPECTED_PROOF);
  });
});

describe("noise NK responder ↔ initiator round trip (pure JS)", () => {
  it("completes handshake, encrypts/decrypts, and verifies auth proof", async () => {
    const initiatorStatic = x25519GenerateKeyPair();
    const responderStatic = x25519GenerateKeyPair();

    const initToResp: Uint8Array[] = [];
    const respToInit: Uint8Array[] = [];

    const initSend = async (b: Uint8Array) => { initToResp.push(b); };
    const initRecv = async () => {
      while (respToInit.length === 0) await new Promise<void>((r) => setTimeout(r, 0));
      return respToInit.shift()!;
    };
    const respSend = async (b: Uint8Array) => { respToInit.push(b); };
    const respRecv = async () => {
      while (initToResp.length === 0) await new Promise<void>((r) => setTimeout(r, 0));
      return initToResp.shift()!;
    };

    const [initSession, respSession] = await Promise.all([
      noiseNKInitiator(responderStatic.publicKey, initSend, initRecv),
      noiseNKResponder(responderStatic.secretKey, responderStatic.publicKey, respSend, respRecv),
    ]);

    // Transport phase: bidirectional encrypt/decrypt.
    const initiatorOut = initSession.encrypt(new TextEncoder().encode("hello from initiator"));
    expect(new TextDecoder().decode(respSession.decrypt(initiatorOut))).toBe("hello from initiator");
    const responderOut = respSession.encrypt(new TextEncoder().encode("hello from responder"));
    expect(new TextDecoder().decode(initSession.decrypt(responderOut))).toBe("hello from responder");

    // Auth round.
    const initiatorEncSend = async (data: Uint8Array) => { initToResp.push(initSession.encrypt(data)); };
    const initiatorEncRecv = async () => initSession.decrypt(await initRecv());
    const responderEncSend = async (data: Uint8Array) => { respToInit.push(respSession.encrypt(data)); };
    const responderEncRecv = async () => respSession.decrypt(await respRecv());

    const [, authResult] = await Promise.all([
      performInitiatorAuth(
        initSession,
        initiatorStatic.secretKey,
        initiatorStatic.publicKey,
        responderStatic.publicKey,
        initiatorEncSend,
        initiatorEncRecv,
      ),
      performResponderAuth(
        respSession,
        responderStatic.secretKey,
        responderStatic.publicKey,
        responderEncSend,
        responderEncRecv,
      ),
    ]);

    expect(bytesToHex(authResult.initiatorPublicKey)).toBe(bytesToHex(initiatorStatic.publicKey));
  });

  it("rejects a forged proof (constant-time MAC compare)", async () => {
    const initiatorStatic = x25519GenerateKeyPair();
    const attackerStatic = x25519GenerateKeyPair();
    const responderStatic = x25519GenerateKeyPair();

    const initToResp: Uint8Array[] = [];
    const respToInit: Uint8Array[] = [];
    const initSend = async (b: Uint8Array) => { initToResp.push(b); };
    const initRecv = async () => {
      while (respToInit.length === 0) await new Promise<void>((r) => setTimeout(r, 0));
      return respToInit.shift()!;
    };
    const respSend = async (b: Uint8Array) => { respToInit.push(b); };
    const respRecv = async () => {
      while (initToResp.length === 0) await new Promise<void>((r) => setTimeout(r, 0));
      return initToResp.shift()!;
    };

    const [initSession, respSession] = await Promise.all([
      noiseNKInitiator(responderStatic.publicKey, initSend, initRecv),
      noiseNKResponder(responderStatic.secretKey, responderStatic.publicKey, respSend, respRecv),
    ]);

    // Initiator lies — claims initiator's identity but signs with attacker's key.
    const initiatorEncSend = async (data: Uint8Array) => { initToResp.push(initSession.encrypt(data)); };
    const initiatorEncRecv = async () => initSession.decrypt(await initRecv());
    const responderEncSend = async (data: Uint8Array) => { respToInit.push(respSession.encrypt(data)); };
    const responderEncRecv = async () => respSession.decrypt(await respRecv());

    const initiatorPromise = performInitiatorAuth(
      initSession,
      attackerStatic.secretKey, // wrong private — used to compute proof
      initiatorStatic.publicKey, // claimed identity
      responderStatic.publicKey,
      initiatorEncSend,
      initiatorEncRecv,
    );
    const responderPromise = performResponderAuth(
      respSession,
      responderStatic.secretKey,
      responderStatic.publicKey,
      responderEncSend,
      responderEncRecv,
    );

    await expect(Promise.all([initiatorPromise, responderPromise])).rejects.toThrow(/proof did not verify/);
  });
});

describe("PCH-01-003b — NK auth proof binds the Noise handshake hash", () => {
  // performInitiatorAuth/performResponderAuth only read session.handshakeHash and
  // the send/recv callbacks, so plaintext channels + minimal fake sessions suffice.
  function plaintextChannels() {
    const initToResp: Uint8Array[] = [];
    const respToInit: Uint8Array[] = [];
    const drain = async (q: Uint8Array[]) => {
      while (q.length === 0) await new Promise<void>((r) => setTimeout(r, 0));
      return q.shift()!;
    };
    return {
      initiatorEncSend: async (d: Uint8Array) => { initToResp.push(d); },
      initiatorEncRecv: async () => drain(respToInit),
      responderEncSend: async (d: Uint8Array) => { respToInit.push(d); },
      responderEncRecv: async () => drain(initToResp),
    };
  }

  it("auth_proof_roundtrips_with_handshake_hash_binding", async () => {
    const initiatorStatic = x25519GenerateKeyPair();
    const responderStatic = x25519GenerateKeyPair();
    const hash = new Uint8Array(32).fill(0x07);
    const ch = plaintextChannels();

    const [, authResult] = await Promise.all([
      performInitiatorAuth(
        { handshakeHash: hash } as never,
        initiatorStatic.secretKey,
        initiatorStatic.publicKey,
        responderStatic.publicKey,
        ch.initiatorEncSend,
        ch.initiatorEncRecv,
      ),
      performResponderAuth(
        { handshakeHash: hash } as never,
        responderStatic.secretKey,
        responderStatic.publicKey,
        ch.responderEncSend,
        ch.responderEncRecv,
      ),
    ]);

    expect(bytesToHex(authResult.initiatorPublicKey)).toBe(bytesToHex(initiatorStatic.publicKey));
  });

  it("auth_proof_rejected_when_handshake_hash_differs", async () => {
    const initiatorStatic = x25519GenerateKeyPair();
    const responderStatic = x25519GenerateKeyPair();
    const ch = plaintextChannels();
    // Same DH/keys/nonce, but the two sides hold DIFFERENT handshake hashes —
    // as a cross-session replay would. The proof must fail to verify.
    const initiatorPromise = performInitiatorAuth(
      { handshakeHash: new Uint8Array(32).fill(0x07) } as never,
      initiatorStatic.secretKey,
      initiatorStatic.publicKey,
      responderStatic.publicKey,
      ch.initiatorEncSend,
      ch.initiatorEncRecv,
    );
    const responderPromise = performResponderAuth(
      { handshakeHash: new Uint8Array(32).fill(0x09) } as never,
      responderStatic.secretKey,
      responderStatic.publicKey,
      ch.responderEncSend,
      ch.responderEncRecv,
    );

    await expect(Promise.all([initiatorPromise, responderPromise])).rejects.toThrow(/proof did not verify/);
  });
});

function hexToBytes(hex: string): Uint8Array {
  const out = new Uint8Array(hex.length / 2);
  for (let i = 0; i < hex.length; i += 2) {
    out[i / 2] = parseInt(hex.slice(i, i + 2), 16);
  }
  return out;
}