// Noise_NX_25519_ChaChaPoly_SHA256 initiator (SPEC.md §4), rev-34 semantics.
// The initiator is anonymous; the responder's static key arrives encrypted in
// message two, which is what server pinning checks.

import { aeadDecrypt, aeadEncrypt, concat, hkdfSha256, randomBytes, sha256, utf8Bytes, x25519, x25519Base } from "./crypto.js";

const PROTOCOL = "Noise_NX_25519_ChaChaPoly_SHA256";  // exactly 32 bytes, so h = name
const PROLOGUE = utf8Bytes("smolmail/1");

// Noise's ChaChaPoly nonce: 4 zero bytes then the counter as u64 LE.
function nonce(n) {
  const out = new Uint8Array(12);
  new DataView(out.buffer).setBigUint64(4, BigInt(n), true);
  return out;
}

// One direction of the post-handshake transport. The key is unique per
// session, so the counter starting at zero is safe.
class CipherState {
  constructor(key) {
    this.key = key;
    this.counter = 0;
  }

  encrypt(plaintext) {
    const sealed = aeadEncrypt(this.key, nonce(this.counter), plaintext, new Uint8Array(0));
    this.counter++;
    return sealed;
  }

  decrypt(sealed) {
    const plaintext = aeadDecrypt(this.key, nonce(this.counter), sealed, new Uint8Array(0));
    this.counter++;
    return plaintext;
  }
}

export class NxInitiator {
  constructor() {
    this.h = utf8Bytes(PROTOCOL);
    this.ck = this.h.slice();
    this.mixHash(PROLOGUE);
    this.key = null;
  }

  mixHash(data) {
    this.h = sha256(concat(this.h, data));
  }

  mixKey(ikm) {
    const okm = hkdfSha256(ikm, this.ck, new Uint8Array(0), 64);
    this.ck = okm.slice(0, 32);
    this.key = okm.slice(32);
  }

  // Message one is just our ephemeral public key. No key is set yet, so the
  // empty payload travels in the clear — and is still mixed into h.
  // `esk` is pinned only by the test vectors, like the spec's fixed-ephemeral
  // envelope.
  writeMessage1(esk) {
    this.esk = esk ?? randomBytes(32);
    this.epk = x25519Base(this.esk);
    this.mixHash(this.epk);
    this.mixHash(new Uint8Array(0));
    return this.epk;
  }

  // Message two: e (plaintext), ee, then the responder's static and the
  // (empty) payload as AEAD ciphertexts chained through h. Each MixKey
  // restarts the nonce at zero.
  readMessage2(message) {
    if (message.length !== 32 + 48 + 16)
      throw new Error(`unexpected NX message length ${message.length}`);
    const re = message.slice(0, 32);
    this.mixHash(re);
    this.mixKey(x25519(this.esk, re));
    this.serverStatic = this.decryptAndHash(message.slice(32, 80));
    this.mixKey(x25519(this.esk, this.serverStatic));  // es
    const payload = this.decryptAndHash(message.slice(80));
    if (payload.length !== 0) throw new Error("unexpected payload in handshake");
    this.handshakeHash = this.h;
    // Split(): two transport keys from the final chaining key, zero-length ikm
    const okm = hkdfSha256(new Uint8Array(0), this.ck, new Uint8Array(0), 64);
    return { send: new CipherState(okm.slice(0, 32)), recv: new CipherState(okm.slice(32)) };
  }

  decryptAndHash(sealed) {
    const plaintext = aeadDecrypt(this.key, nonce(0), sealed, this.h);
    this.mixHash(sealed);
    return plaintext;
  }
}