/**
 * Noise NK initiator for Android / Capacitor.
 * Pattern: Noise_NK_25519_ChaChaPoly_BLAKE2s
 *
 * The native Rust host runs `snow` to drive the same pattern; this JS
 * implementation produces wire-identical bytes so a Capacitor WebView client
 * can connect to the desktop host over WebSocket.
 *
 * MUST stay in lockstep with the native side. The cross-check test in
 * `noise-js.test.ts` pins exact byte outputs against fixed inputs so a
 * regression on either side fails CI before any device pair gets wedged.
 */

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

const PROTOCOL_NAME = "Noise_NK_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;
  }
  const sha = new SHA256();
  sha.update(nameBytes);
  return sha.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] {
  const out = new HKDF(SHA256, ikm, ck).expand(64);
  return [out.slice(0, 32), out.slice(32, 64)];
}

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

/**
 * 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> {
  let h: Uint8Array = initHash();
  let ck: Uint8Array = h.slice();

  // NK pre-message: mix responder static public key.
  h = mixHash(h, responderStaticPublicKey);

  const ephemeral = generateKeyPair();

  // --- Message 1: e, es ---
  h = mixHash(h, ephemeral.publicKey);
  const es = sharedKey(ephemeral.secretKey, responderStaticPublicKey);
  const [ck1, k1] = hkdf2(ck, es);
  ck = ck1;

  const cipher1 = new ChaCha20Poly1305(k1);
  const nonce1 = new Uint8Array(12);
  const encrypted1 = cipher1.seal(nonce1, new Uint8Array(0), h);
  h = mixHash(h, encrypted1);

  const msg1 = new Uint8Array(32 + encrypted1.length);
  msg1.set(ephemeral.publicKey, 0);
  msg1.set(encrypted1, 32);
  await send(msg1);

  // --- Message 2: e, ee ---
  const msg2 = await receive();
  if (msg2.length < 48) throw new Error("Noise NK: msg2 too short");

  const remoteEphemeralPublic = msg2.slice(0, 32);
  const encrypted2 = msg2.slice(32);

  h = mixHash(h, remoteEphemeralPublic);
  const ee = sharedKey(ephemeral.secretKey, remoteEphemeralPublic);
  const [ck2, k2] = hkdf2(ck, ee);
  ck = ck2;

  const cipher2 = new ChaCha20Poly1305(k2);
  const nonce2 = new Uint8Array(12);
  const decrypted2 = cipher2.open(nonce2, encrypted2, h);
  if (!decrypted2) {
    throw new Error("Noise NK: msg2 auth failed — wrong responder key or MITM");
  }
  h = mixHash(h, encrypted2);

  // --- Transport phase ---
  const [sendKey, recvKey] = hkdf2(ck, new Uint8Array(0));
  const handshakeHash = h.slice();

  let sendNonce = 0n;
  let recvNonce = 0n;

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

  const sendCipher = new ChaCha20Poly1305(sendKey);
  const recvCipher = new ChaCha20Poly1305(recvKey);

  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 NK: decryption failed");
      return result;
    },
    handshakeHash,
  };
}

/**
 * Compute the auth proof byte-for-byte identically to the native host.
 * See `native/peach-sync-host/src/auth.rs::proof_for`.
 *
 * 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
 * )
 *
 * The initiator and responder both pass the same DH peer key (their PEER's
 * public key) and the same initiator_public/responder_public ordering so the
 * MAC matches. `mac_for` in auth.rs implements both directions; 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,
): Uint8Array {
  if (nonce.length !== 32) throw new Error("auth nonce must be 32 bytes");

  const dh = sharedKey(initiatorPrivate, responderPublic);
  const mac = new HMAC(SHA256, dh);
  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).
 */
export async function performInitiatorAuth(
  session: NoiseNKSession,
  ourPrivateKey: Uint8Array,
  ourPublicKey: Uint8Array,
  responderPublicKey: Uint8Array,
  sendEncryptedJson: (data: Uint8Array) => Promise<void>,
  receiveEncryptedJson: () => Promise<Uint8Array>,
): Promise<void> {
  void session; // hash bound into MAC via the DH key only — kept for future fields

  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);
  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 };