/**
 * Noise NK initiator for browser-to-browser sync.
 * Pattern: Noise_NK_25519_ChaChaPoly_BLAKE2s
 *
 * The protocol state machine lives in the canonical Rust `snow` core and is
 * compiled to WASM here and JNI on Android. This file keeps the stable
 * TypeScript orchestration API plus the post-handshake proof helpers.
 */

import { sharedKey } from "@stablelib/x25519";
import { BLAKE2s } from "@stablelib/blake2s";
import { SHA256 } from "@stablelib/sha256";
import { HMAC } from "@stablelib/hmac";
import {
  rustNoiseKkInitiator,
  rustNoiseKkResponder,
  rustNoiseNkInitiator,
  rustNoiseNkResponder,
} from "./noise-core";

const PROTOCOL_NAME = "Noise_NK_25519_ChaChaPoly_BLAKE2s";

// libhydrogen's contributory-behaviour blocklist: the nine small-order X25519
// u-coordinates. Mirrors X25519_SMALL_ORDER_PUBKEYS in crypto-core (PB-26-014)
// so the TS-side DH boundaries fail before the DH, not just the WASM core.
const X25519_SMALL_ORDER_PUBKEYS = new Set([
  "00".repeat(32),
  "01" + "00".repeat(31),
  "e0eb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b880",
  "5f9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f11d7",
  "ec" + "ff".repeat(30) + "7f",
  "ed" + "ff".repeat(30) + "7f",
  "ee" + "ff".repeat(30) + "7f",
  "cdeb7a7c3b41b8ae1656e3faf19fc46ada098deb9c32b1fd866205165f49b880",
  "4c9c95bca3508c24b1d0b1559c83ef5b04445cc4581c8e86d8224eddd09f11d7",
]);

export function isSmallOrderPubkey(publicKey: Uint8Array): boolean {
  return publicKey.length === 32 && X25519_SMALL_ORDER_PUBKEYS.has(bytesToHex(publicKey));
}

function initProtocolHash(protocolName: string): Uint8Array {
  const nameBytes = new TextEncoder().encode(protocolName);
  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.
  return new BLAKE2s(32).update(nameBytes).digest();
}

export function initHash(): Uint8Array {
  return initProtocolHash(PROTOCOL_NAME);
}

export 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 blocks and BLAKE2s's restoreState()
  // never clears its `finished` flag, so `new HKDF(BLAKE2s, …)` throws. Run
  // standard HKDF-Expand (RFC 5869, empty info) with a FRESH HMAC-BLAKE2s per
  // block — byte-identical to BouncyCastle `HKDFBytesGenerator(Blake2sDigest)`
  // on peach-android. NOTE: the HMAC-SHA256 post-handshake auth proof is a
  // SEPARATE construction, intentionally unchanged (so the SHA256 import stays).
  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];
}

export interface NoiseNKSession {
  encrypt(plaintext: Uint8Array): Uint8Array;
  decrypt(ciphertext: Uint8Array): Uint8Array;
  readonly handshakeHash: Uint8Array;
  rekeyOutgoing(): void;
  rekeyIncoming(): void;
  close(): void;
}

/**
 * Standard Noise KK initiator for reconnects.
 *
 * Pre-messages: -> s, <- s
 * Message 1:    -> e, es, ss
 * Message 2:    <- e, ee, se
 *
 * Both static keys are local inputs. There is no negotiation and no fallback
 * to NK: a mismatch hard-fails this handshake.
 */
export async function noiseKKInitiator(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  responderStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<NoiseNKSession> {
  if (isSmallOrderPubkey(responderStaticPublicKey)) {
    throw new Error("noise: small-order X25519 responder static key");
  }
  return rustNoiseKkInitiator(
    ourStaticPrivateKey,
    ourStaticPublicKey,
    responderStaticPublicKey,
    send,
    receive,
  );
}

/**
 * Standard Noise KK responder. The expected initiator key must come from the
 * local paired-device registry before this function is called.
 */
export async function noiseKKResponder(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  expectedInitiatorStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<NoiseNKSession> {
  return rustNoiseKkResponder(
    ourStaticPrivateKey,
    ourStaticPublicKey,
    expectedInitiatorStaticPublicKey,
    send,
    receive,
  );
}

/**
 * Run the initiator side of the Noise NK handshake. The caller provides
 * `send` / `receive` for the two transport-side handshake messages and gets
 * back a session object that encrypts/decrypts subsequent transport frames.
 */
export async function noiseNKInitiator(
  responderStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<NoiseNKSession> {
  if (isSmallOrderPubkey(responderStaticPublicKey)) {
    throw new Error("noise: small-order X25519 responder static key");
  }
  return rustNoiseNkInitiator(
    responderStaticPublicKey,
    send,
    receive,
  );
}

/**
 * Compute the post-handshake auth proof.
 *
 * proof = HMAC-SHA256(
 *   key   = X25519(local_priv, remote_pub),  // 32-byte DH shared
 *   data  = "peach-sync-v1-initiator-auth"
 *        || initiator_public
 *        || responder_public
 *        || nonce_32
 * )
 *
 * Both sides pass the same DH peer key (their PEER's public key) and the same
 * initiator_public/responder_public ordering so the MAC matches. This helper
 * targets the initiator-side: `local_priv` is the initiator's private key,
 * `remote_pub` is the responder's public key.
 */
const AUTH_DOMAIN = new TextEncoder().encode("peach-sync-v1-initiator-auth");

export function computeAuthProof(
  initiatorPrivate: Uint8Array,
  initiatorPublic: Uint8Array,
  responderPublic: Uint8Array,
  nonce: Uint8Array,
  handshakeHash: Uint8Array,
): Uint8Array {
  if (nonce.length !== 32) throw new Error("auth nonce must be 32 bytes");
  if (handshakeHash.length !== 32) throw new Error("handshake hash must be 32 bytes");
  if (isSmallOrderPubkey(responderPublic)) {
    throw new Error("small-order X25519 responder public key");
  }

  const dh = sharedKey(initiatorPrivate, responderPublic, true);
  const mac = new HMAC(SHA256, dh);
  // PCH-01-003b: bind the proof to THIS Noise session so it cannot be replayed
  // onto a different handshake. handshakeHash is prepended; both sides hold the
  // identical value and feed it in the same position.
  mac.update(handshakeHash);
  mac.update(AUTH_DOMAIN);
  mac.update(initiatorPublic);
  mac.update(responderPublic);
  mac.update(nonce);
  return mac.digest();
}

interface AuthChallengeFrame {
  type: "CHALLENGE";
  nonce_hex: string;
  responder_public_key_hex: string;
}

interface AuthProofFrame {
  type: "PROOF";
  initiator_public_key_hex: string;
  proof_hex: string;
}

function bytesToHex(bytes: Uint8Array): string {
  return Array.from(bytes, (b) => b.toString(16).padStart(2, "0")).join("");
}

function hexToBytes(hex: string): Uint8Array {
  if (hex.length % 2 !== 0) throw new Error("invalid hex");
  const bytes = new Uint8Array(hex.length / 2);
  for (let i = 0; i < hex.length; i += 2) {
    bytes[i / 2] = parseInt(hex.slice(i, i + 2), 16);
  }
  return bytes;
}

/**
 * Drive the initiator side of the post-handshake challenge/proof round. The
 * responder sends a `Challenge` JSON frame; we compute the HMAC and reply
 * with a `Proof` frame. Both frames are encrypted via the post-handshake
 * Noise transport (caller owns the `sendEncrypted` / `receiveEncrypted` pair).
 */
/**
 * Run the responder side of the Noise NK handshake. Symmetric with
 * `noiseNKInitiator` — the caller provides `send` / `receive` for the two
 * handshake messages and gets back a session for transport-phase frames.
 *
 * The Rust core runs the same state machine as the browser initiator and the
 * Android JNI adapter.
 */
export async function noiseNKResponder(
  ourStaticPrivateKey: Uint8Array,
  ourStaticPublicKey: Uint8Array,
  send: (msg: Uint8Array) => Promise<void>,
  receive: () => Promise<Uint8Array>,
): Promise<NoiseNKSession> {
  return rustNoiseNkResponder(
    ourStaticPrivateKey,
    ourStaticPublicKey,
    send,
    receive,
  );
}

/**
 * Constant-time 32-byte equality. The MAC comparison in `verifyAuthProof`
 * goes through this so a same-LAN attacker can't time the HMAC compare.
 */
function constantTimeEq(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  let diff = 0;
  for (let i = 0; i < a.length; i++) diff |= a[i]! ^ b[i]!;
  return diff === 0;
}

/**
 * Drive the responder side of the post-handshake challenge/proof round. We
 * generate a random 32-byte nonce, send `Challenge`, receive `Proof`, and
 * verify. Returns the initiator's authenticated public key on success.
 *
 * The MAC key is X25519(our_priv, initiator_pub) and the data is the same
 * domain-separated transcript both sides feed into HMAC-SHA256.
 */
export async function performResponderAuth(
  session: NoiseNKSession,
  ourPrivateKey: Uint8Array,
  ourPublicKey: Uint8Array,
  sendEncryptedJson: (data: Uint8Array) => Promise<void>,
  receiveEncryptedJson: () => Promise<Uint8Array>,
): Promise<{ initiatorPublicKey: Uint8Array }> {
  // PCH-01-003b: bind the verifier's recomputed MAC to this Noise session.

  const nonce = new Uint8Array(32);
  crypto.getRandomValues(nonce);

  const challenge: AuthChallengeFrame = {
    type: "CHALLENGE",
    nonce_hex: bytesToHex(nonce),
    responder_public_key_hex: bytesToHex(ourPublicKey),
  };
  await sendEncryptedJson(new TextEncoder().encode(JSON.stringify(challenge)));

  const proofBytes = await receiveEncryptedJson();
  const proofFrame = JSON.parse(new TextDecoder().decode(proofBytes)) as AuthProofFrame;
  if (proofFrame.type !== "PROOF") {
    throw new Error("expected PROOF frame");
  }

  const initiatorPublicKey = hexToBytes(proofFrame.initiator_public_key_hex);
  const claimedProof = hexToBytes(proofFrame.proof_hex);
  if (initiatorPublicKey.length !== 32) {
    throw new Error("initiator public key must be 32 bytes");
  }
  if (claimedProof.length !== 32) {
    throw new Error("proof must be 32 bytes");
  }
  if (isSmallOrderPubkey(initiatorPublicKey)) {
    throw new Error("small-order X25519 initiator public key");
  }

  // mac_for(our_priv, initiator_pub_for_dh, initiator_pub, our_pub, nonce)
  const dh = sharedKey(ourPrivateKey, initiatorPublicKey, true);
  const mac = new HMAC(SHA256, dh);
  // PCH-01-003b: must mirror computeAuthProof — handshakeHash first, then the
  // existing domain/keys/nonce transcript. Same value, same position on both ends.
  mac.update(session.handshakeHash);
  mac.update(AUTH_DOMAIN);
  mac.update(initiatorPublicKey);
  mac.update(ourPublicKey);
  mac.update(nonce);
  const expected = mac.digest();

  if (!constantTimeEq(expected, claimedProof)) {
    throw new Error("auth proof did not verify");
  }
  return { initiatorPublicKey };
}

export async function performInitiatorAuth(
  session: NoiseNKSession,
  ourPrivateKey: Uint8Array,
  ourPublicKey: Uint8Array,
  responderPublicKey: Uint8Array,
  sendEncryptedJson: (data: Uint8Array) => Promise<void>,
  receiveEncryptedJson: () => Promise<Uint8Array>,
): Promise<void> {
  // PCH-01-003b: the Noise handshake hash is now mixed into the proof, binding
  // it to this exact session (was previously bound via the DH key only).

  const challengeBytes = await receiveEncryptedJson();
  const challenge = JSON.parse(new TextDecoder().decode(challengeBytes)) as AuthChallengeFrame;
  if (challenge.type !== "CHALLENGE") {
    throw new Error("expected CHALLENGE frame");
  }
  const nonce = hexToBytes(challenge.nonce_hex);
  const responderKeyFromChallenge = hexToBytes(challenge.responder_public_key_hex);
  // Defence against an off-key challenge: the responder's claimed key MUST be
  // the one we Noise-NK-handshook against, otherwise an active MITM that
  // managed to substitute a different key is rerouting the auth round.
  if (
    responderKeyFromChallenge.length !== responderPublicKey.length ||
    !responderKeyFromChallenge.every((b, i) => b === responderPublicKey[i])
  ) {
    throw new Error("responder key in challenge does not match Noise handshake");
  }

  const proof = computeAuthProof(
    ourPrivateKey,
    ourPublicKey,
    responderPublicKey,
    nonce,
    session.handshakeHash,
  );
  const proofFrame: AuthProofFrame = {
    type: "PROOF",
    initiator_public_key_hex: bytesToHex(ourPublicKey),
    proof_hex: bytesToHex(proof),
  };
  await sendEncryptedJson(new TextEncoder().encode(JSON.stringify(proofFrame)));
}

export { bytesToHex, hexToBytes };