// Noise XX (mutual auth) — used for every Phase-2 WebSocket connection
// after the audio chirp pairing has stamped each side into the other's
// paired-device registry. Both peers hold a long-term X25519 static key
// (the sync identity); XX is the canonical pattern when both parties want
// to authenticate each other.
//
// Pattern: Noise_XX_25519_ChaChaPoly_BLAKE2s — same primitives as the
// NK pattern in `noise-js.ts`, so the helpers are duplicated minimally
// (initHash / mixHash / hkdf2) to keep this file self-contained.
//
//   -> e
//   <- e, ee, s, es
//   -> s, se
//
// After msg 3, the symmetric `hkdf2(ck, "")` splits into two ChaCha20-
// Poly1305 transport keys, one per direction. Counter nonces start at 0
// and increment per frame.

import { generateKeyPair, sharedKey } from "@stablelib/x25519";
import { ChaCha20Poly1305 } from "@stablelib/chacha20poly1305";
import { BLAKE2s } from "@stablelib/blake2s";
import { HMAC } from "@stablelib/hmac";

const PROTOCOL_NAME = "Noise_XX_25519_ChaChaPoly_BLAKE2s";

function initHash(): Uint8Array {
  const nameBytes = new TextEncoder().encode(PROTOCOL_NAME);
  if (nameBytes.length <= 32) {
    const h = new Uint8Array(32);
    h.set(nameBytes);
    return h;
  }
  // PCH-01-002: the named pattern hash is BLAKE2s; the protocol-name init hash
  // MUST use it too (Noise spec §5.2), not SHA-256. BLAKE2s digest is 32 bytes,
  // matching mixHash output width.
  return new BLAKE2s(32).update(nameBytes).digest();
}

function mixHash(h: Uint8Array, data: Uint8Array): Uint8Array {
  const blake = new BLAKE2s(32);
  blake.update(h);
  blake.update(data);
  return blake.digest();
}

function hkdf2(ck: Uint8Array, ikm: Uint8Array): [Uint8Array, Uint8Array] {
  // PCH-01-002: HKDF must run over the named hash (BLAKE2s), not SHA-256.
  // @stablelib's HKDF reuses one HMAC across output blocks and BLAKE2s's
  // restoreState() never clears its `finished` flag, so `new HKDF(BLAKE2s, …)`
  // throws "can't update because hash was finished". Run standard HKDF-Expand
  // (RFC 5869, empty info) with a FRESH HMAC-BLAKE2s per block instead — this is
  // byte-identical to BouncyCastle `HKDFBytesGenerator(Blake2sDigest(256))` on
  // the peach-android side, so the two stay in wire lockstep.
  const prk = new HMAC(BLAKE2s, ck).update(ikm).digest(); // Extract (salt=ck)
  const k1 = new HMAC(BLAKE2s, prk).update(Uint8Array.of(1)).digest();
  const k2 = new HMAC(BLAKE2s, prk).update(k1).update(Uint8Array.of(2)).digest();
  return [k1, k2];
}

function nonceBytes(n: bigint): Uint8Array {
  const buf = new Uint8Array(12);
  new DataView(buf.buffer).setBigUint64(4, n, true);
  return buf;
}

export interface NoiseXXSession {
  encrypt(plaintext: Uint8Array): Uint8Array;
  decrypt(ciphertext: Uint8Array): Uint8Array;
  remoteStaticPublicKey: Uint8Array;
  handshakeHash: Uint8Array;
}

/**
 * Drive the XX initiator. The responder's static key is learned in
 * message 2; we hand it back so the caller can verify it against the
 * paired-device registry.
 */
export async function noiseXXInitiator(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
  // PCH-01-003: called once the responder's static is authenticated (end of
  // msg2) and BEFORE msg3 — which carries OUR static — is sent. Throwing here
  // aborts the handshake without disclosing our long-term identity to an
  // unexpected peer. Honest callers pass a pin check; omitting it preserves the
  // old behaviour.
  verifyRemoteStatic?: (remoteStaticPublicKey: Uint8Array) => void,
): Promise<NoiseXXSession> {
  let h: Uint8Array = initHash();
  let ck: Uint8Array = h.slice();

  // --- Message 1: -> e ---
  const e = generateKeyPair();
  h = mixHash(h, e.publicKey);
  await send(e.publicKey);

  // --- Message 2: <- e, ee, s, es ---
  const msg2 = await receive();
  if (msg2.length < 32 + 32 + 16 + 16) {
    throw new Error("Noise XX: msg2 too short");
  }
  const remoteEphemeral = msg2.slice(0, 32);
  const encryptedStatic = msg2.slice(32, 32 + 32 + 16);
  const encryptedPayload = msg2.slice(32 + 32 + 16);

  h = mixHash(h, remoteEphemeral);
  const ee = sharedKey(e.secretKey, remoteEphemeral);
  const [ck1, k1] = hkdf2(ck, ee);
  ck = ck1;

  const cipherStatic = new ChaCha20Poly1305(k1);
  const remoteStaticPlain = cipherStatic.open(nonceBytes(0n), encryptedStatic, h);
  if (!remoteStaticPlain || remoteStaticPlain.length !== 32) {
    throw new Error("Noise XX: msg2 static auth failed");
  }
  h = mixHash(h, encryptedStatic);

  const es = sharedKey(e.secretKey, remoteStaticPlain);
  const [ck2, k2] = hkdf2(ck, es);
  ck = ck2;

  const cipherPayload = new ChaCha20Poly1305(k2);
  const payloadPlain = cipherPayload.open(nonceBytes(0n), encryptedPayload, h);
  if (!payloadPlain) {
    throw new Error("Noise XX: msg2 payload auth failed");
  }
  h = mixHash(h, encryptedPayload);

  // PCH-01-003: pin the responder static BEFORE we transmit our own static in
  // msg3. `remoteStaticPlain` was authenticated above (msg2). If the caller's
  // check throws, we return without sending msg3 — our identity stays private.
  if (verifyRemoteStatic) verifyRemoteStatic(remoteStaticPlain);

  // --- Message 3: -> s, se ---
  const cipherOurStatic = new ChaCha20Poly1305(k2);
  const encryptedOurStatic = cipherOurStatic.seal(nonceBytes(1n), ourStaticPublicKey, h);
  h = mixHash(h, encryptedOurStatic);

  const se = sharedKey(ourStaticPrivateKey, remoteEphemeral);
  const [ck3, k3] = hkdf2(ck, se);
  ck = ck3;

  const cipherFinal = new ChaCha20Poly1305(k3);
  const encryptedFinalPayload = cipherFinal.seal(nonceBytes(0n), new Uint8Array(0), h);
  h = mixHash(h, encryptedFinalPayload);

  const msg3 = new Uint8Array(encryptedOurStatic.length + encryptedFinalPayload.length);
  msg3.set(encryptedOurStatic, 0);
  msg3.set(encryptedFinalPayload, encryptedOurStatic.length);
  await send(msg3);

  // --- Transport keys: split (k_init_send, k_resp_send) ---
  const [initSend, respSend] = hkdf2(ck, new Uint8Array(0));
  const sendCipher = new ChaCha20Poly1305(initSend);
  const recvCipher = new ChaCha20Poly1305(respSend);
  let sendNonce = 0n;
  let recvNonce = 0n;

  return {
    encrypt(plaintext) {
      return sendCipher.seal(nonceBytes(sendNonce++), plaintext, new Uint8Array(0));
    },
    decrypt(ciphertext) {
      const result = recvCipher.open(nonceBytes(recvNonce++), ciphertext, new Uint8Array(0));
      if (!result) throw new Error("Noise XX: decryption failed");
      return result;
    },
    remoteStaticPublicKey: remoteStaticPlain,
    handshakeHash: h.slice(),
  };
}

/**
 * Drive the XX responder. The initiator's static key is learned in
 * message 3; we hand it back on success.
 */
export async function noiseXXResponder(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<NoiseXXSession> {
  let h: Uint8Array = initHash();
  let ck: Uint8Array = h.slice();

  // --- Message 1: <- e ---
  const msg1 = await receive();
  if (msg1.length < 32) throw new Error("Noise XX: msg1 too short");
  const remoteEphemeral = msg1.slice(0, 32);
  h = mixHash(h, remoteEphemeral);

  // --- Message 2: -> e, ee, s, es ---
  const e = generateKeyPair();
  h = mixHash(h, e.publicKey);
  const ee = sharedKey(e.secretKey, remoteEphemeral);
  const [ck1, k1] = hkdf2(ck, ee);
  ck = ck1;

  const cipherStatic = new ChaCha20Poly1305(k1);
  const encryptedOurStatic = cipherStatic.seal(nonceBytes(0n), ourStaticPublicKey, h);
  h = mixHash(h, encryptedOurStatic);

  const es = sharedKey(ourStaticPrivateKey, remoteEphemeral);
  const [ck2, k2] = hkdf2(ck, es);
  ck = ck2;

  const cipherPayload = new ChaCha20Poly1305(k2);
  const encryptedPayload = cipherPayload.seal(nonceBytes(0n), new Uint8Array(0), h);
  h = mixHash(h, encryptedPayload);

  const msg2 = new Uint8Array(e.publicKey.length + encryptedOurStatic.length + encryptedPayload.length);
  msg2.set(e.publicKey, 0);
  msg2.set(encryptedOurStatic, e.publicKey.length);
  msg2.set(encryptedPayload, e.publicKey.length + encryptedOurStatic.length);
  await send(msg2);

  // --- Message 3: <- s, se ---
  const msg3 = await receive();
  if (msg3.length < 32 + 16 + 16) throw new Error("Noise XX: msg3 too short");
  const encryptedRemoteStatic = msg3.slice(0, 32 + 16);
  const encryptedFinalPayload = msg3.slice(32 + 16);

  const cipherTheirStatic = new ChaCha20Poly1305(k2);
  const remoteStaticPlain = cipherTheirStatic.open(nonceBytes(1n), encryptedRemoteStatic, h);
  if (!remoteStaticPlain || remoteStaticPlain.length !== 32) {
    throw new Error("Noise XX: msg3 static auth failed");
  }
  h = mixHash(h, encryptedRemoteStatic);

  const se = sharedKey(e.secretKey, remoteStaticPlain);
  const [ck3, k3] = hkdf2(ck, se);
  ck = ck3;

  const cipherFinal = new ChaCha20Poly1305(k3);
  const finalPayloadPlain = cipherFinal.open(nonceBytes(0n), encryptedFinalPayload, h);
  if (!finalPayloadPlain) {
    throw new Error("Noise XX: msg3 final payload auth failed");
  }
  h = mixHash(h, encryptedFinalPayload);

  const [initSend, respSend] = hkdf2(ck, new Uint8Array(0));
  const sendCipher = new ChaCha20Poly1305(respSend);
  const recvCipher = new ChaCha20Poly1305(initSend);
  let sendNonce = 0n;
  let recvNonce = 0n;

  return {
    encrypt(plaintext) {
      return sendCipher.seal(nonceBytes(sendNonce++), plaintext, new Uint8Array(0));
    },
    decrypt(ciphertext) {
      const result = recvCipher.open(nonceBytes(recvNonce++), ciphertext, new Uint8Array(0));
      if (!result) throw new Error("Noise XX: decryption failed");
      return result;
    },
    remoteStaticPublicKey: remoteStaticPlain,
    handshakeHash: h.slice(),
  };
}