// 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 "dart:typed_data";

import "package:smol_mail/smol/crypto.dart";
import "package:smol_mail/smol/errors.dart";

const String _protocol = "Noise_NX_25519_ChaChaPoly_SHA256"; // exactly 32 bytes, so h = name

Uint8List _prologue() => utf8Bytes("smolmail/1");

// Noise's ChaChaPoly nonce: 4 zero bytes then the counter as u64 LE.
Uint8List _nonce(int n) {
  final out = Uint8List(12);
  ByteData.view(out.buffer).setUint64(4, n, Endian.little);
  return out;
}

/// One direction of the post-handshake transport; tests substitute a
/// passthrough so framing guards can be exercised without crypto.
abstract class SessionCipher {
  Uint8List encrypt(Uint8List plaintext);

  Uint8List decrypt(Uint8List sealed);
}

// The key is unique per session, so the counter starting at zero is safe.
class CipherState implements SessionCipher {
  final Uint8List key;
  int counter = 0;

  CipherState(this.key);

  @override
  Uint8List encrypt(Uint8List plaintext) {
    final sealed = aeadEncrypt(key, _nonce(counter), plaintext, Uint8List(0));
    counter++;
    return sealed;
  }

  @override
  Uint8List decrypt(Uint8List sealed) {
    final plaintext = aeadDecrypt(key, _nonce(counter), sealed, Uint8List(0));
    counter++;
    return plaintext;
  }
}

class NxResult {
  final CipherState send, recv;
  final Uint8List serverStatic;
  final Uint8List handshakeHash;

  const NxResult(this.send, this.recv, this.serverStatic, this.handshakeHash);
}

class NxInitiator {
  late Uint8List h;
  late Uint8List ck;
  Uint8List? key;
  late Uint8List esk;
  late Uint8List epk;

  NxInitiator() {
    h = utf8Bytes(_protocol);
    ck = Uint8List.fromList(h);
    mixHash(_prologue());
  }

  void mixHash(Uint8List data) {
    h = sha256(concat([h, data]));
  }

  void mixKey(Uint8List ikm) {
    final okm = hkdfSha256(ikm, ck, Uint8List(0), 64);
    ck = okm.sublist(0, 32);
    key = okm.sublist(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.
  Uint8List writeMessage1([Uint8List? esk]) {
    this.esk = esk ?? randomBytes(32);
    epk = x25519Base(this.esk);
    mixHash(epk);
    mixHash(Uint8List(0));
    return Uint8List.fromList(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.
  NxResult readMessage2(Uint8List message) {
    if (message.length != 32 + 48 + 16) {
      throw SmolError("unexpected NX message length ${message.length}");
    }
    final re = message.sublist(0, 32);
    mixHash(re);
    mixKey(x25519(esk, re));
    final serverStatic = decryptAndHash(message.sublist(32, 80));
    mixKey(x25519(esk, serverStatic)); // es
    final payload = decryptAndHash(message.sublist(80));
    if (payload.isNotEmpty) throw const SmolError("unexpected payload in handshake");
    final handshakeHash = h;
    // Split(): two transport keys from the final chaining key, zero-length ikm
    final okm = hkdfSha256(Uint8List(0), ck, Uint8List(0), 64);
    return NxResult(CipherState(okm.sublist(0, 32)), CipherState(okm.sublist(32)),
        serverStatic, handshakeHash);
  }

  Uint8List decryptAndHash(Uint8List sealed) {
    final plaintext = aeadDecrypt(key!, _nonce(0), sealed, h);
    mixHash(sealed);
    return plaintext;
  }
}