// Smol Mail protocol, version 1 (../smolmail SPEC.md): addresses, sealed and
// signed envelopes, body frontmatter, key rotation, and the framed request
// and response bodies of the five operations.

import "dart:math";
import "dart:typed_data";

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

const defaultPort = 1961;
const keyLen = 32, sigLen = 64, certLen = 136, idLen = 16;
const maxFrame = 1 << 20, noisePayload = 65535 - 16, padTo = 1024;
const envelopeHeader = 69, payloadHeader = 45, maxChain = 16;
const _frontmatterMax = 4096, _frontmatterKeys = 64;

const opAuth = 0x00, opResolve = 0x01, opSend = 0x02, opFetch = 0x03,
    opDelete = 0x04, opRegister = 0x05;

const _statusNames = {
  0: "ok", 1: "malformed", 2: "bad version", 3: "unknown user",
  4: "auth required", 5: "auth failed", 6: "quota exceeded", 7: "too large",
  8: "rate limited", 9: "not permitted", 10: "internal error",
};

String statusName(int status) => _statusNames[status] ?? "$status";

final _label = (
  auth: utf8Bytes("smolmail/1 auth"),
  seal: utf8Bytes("smolmail/1 seal"),
  msg: utf8Bytes("smolmail/1 msg"),
  id: utf8Bytes("smolmail/1 id"),
  rotate: utf8Bytes("smolmail/1 rotate"),
);

// --- encoding helpers ---------------------------------------------------------

const _b32 = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";

String b32encode(List<int> bytes) {
  var out = "";
  var value = 0, bits = 0;
  for (final b in bytes) {
    value = (value << 8) | b;
    bits += 8;
    while (bits >= 5) {
      bits -= 5;
      out += _b32[(value >>> bits) & 31];
    }
  }
  if (bits > 0) out += _b32[(value << (5 - bits)) & 31];
  return out.toLowerCase();
}

Uint8List b32decode(String text) {
  final out = <int>[];
  var value = 0, bits = 0;
  final clean = text.trim().toUpperCase().replaceAll(RegExp(r"=+$"), "");
  for (final ch in clean.split("")) {
    final idx = _b32.indexOf(ch);
    if (idx < 0) throw SmolError("invalid base32 character '$ch'");
    value = (value << 5) | idx;
    bits += 5;
    if (bits >= 8) {
      bits -= 8;
      out.add((value >>> bits) & 0xff);
    }
  }
  return Uint8List.fromList(out);
}

// §3: the first 20 base32 characters of the identity, in groups of four.
String fingerprint(Uint8List identity) {
  final s = b32encode(identity).substring(0, 20);
  return RegExp(".{4}").allMatches(s).map((m) => m[0]).join(" ");
}

Uint8List u16be(int n) => Uint8List.fromList([(n >> 8) & 0xff, n & 0xff]);

Uint8List u32be(int n) => Uint8List.fromList(
    [(n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff]);

// Dart 3.2's ByteData has no setBigInt, so write big-endian manually.
Uint8List i64be(BigInt n) {
  final out = Uint8List(8);
  for (var i = 0; i < 8; i++) {
    out[i] = ((n >> (8 * (7 - i))) & BigInt.from(0xff)).toInt();
  }
  return out;
}

int nowSeconds() => DateTime.now().millisecondsSinceEpoch ~/ 1000;

// Fail-closed reader; every parse raises rather than reading past the end.
class Reader {
  final Uint8List buf;
  int pos = 0;

  Reader(this.buf);

  Uint8List take(int n) {
    if (n < 0 || pos + n > buf.length) throw const SmolError("truncated message");
    final out = buf.sublist(pos, pos + n);
    pos += n;
    return out;
  }

  int u8() => take(1)[0];

  int u16() {
    final b = take(2);
    return (b[0] << 8) | b[1];
  }

  int u32() => ByteData.view(take(4).buffer).getUint32(0);

  int i64() => ByteData.view(take(8).buffer).getInt64(0);

  int get left => buf.length - pos;
}

// --- identity -----------------------------------------------------------------

// §2: an Ed25519 keypair with the X25519 agreement keys derived from it.
class SmolIdentity {
  final Uint8List seed;
  final Uint8List publicKey;

  const SmolIdentity(this.seed, this.publicKey);
}

SmolIdentity identityFromSeed(Uint8List seed) {
  if (seed.length != keyLen) {
    throw const SmolError("identity seed must be $keyLen bytes");
  }
  return SmolIdentity(seed, ed25519PublicKey(seed));
}

SmolIdentity newIdentity() => identityFromSeed(randomBytes(keyLen));

// --- addressing (§3) -----------------------------------------------------------

final _address =
    RegExp(r"^(?<user>[a-z0-9._-]{1,63})@(?<host>[^/:]+)(?::(?<port>\d+))?$");

class SmolAddress {
  final String user;
  final String host;
  final int port;

  /// The key carried by a `smol://` address; null for short addresses.
  final Uint8List? identity;

  const SmolAddress(this.user, this.host, this.port, this.identity);

  String get short =>
      "$user@$host${port == defaultPort ? "" : ":$port"}";

  String uri(Uint8List key) =>
      "smol://$user@$host${port == defaultPort ? "" : ":$port"}/${b32encode(key)}";
}

SmolAddress parseAddress(String text) {
  text = text.trim();
  Uint8List? identity;
  if (text.startsWith("smol://")) {
    final rest = text.substring("smol://".length);
    final slash = rest.lastIndexOf("/");
    if (slash < 0) throw SmolError("$text: smol:// address carries no key");
    identity = b32decode(rest.substring(slash + 1));
    if (identity.length != keyLen) {
      throw SmolError(
          "$text: key is ${identity.length} bytes, expected $keyLen");
    }
    text = rest.substring(0, slash);
  }
  final m = _address.firstMatch(text.toLowerCase());
  if (m == null) throw SmolError("'$text' is not a valid address");
  final user = m.namedGroup("user")!;
  final host = m.namedGroup("host")!;
  if ("._-".contains(user[0]) || "._-".contains(user[user.length - 1])) {
    throw SmolError("$user may not begin or end with a separator");
  }
  final portText = m.namedGroup("port");
  final port = portText != null ? int.parse(portText) : defaultPort;
  return SmolAddress(user, host, port, identity);
}

// --- message format (§5) -------------------------------------------------------

Uint8List messageId(List<int> envelope) =>
    sha256(concat([_label.id, envelope])).sublist(0, idLen);

class OpenedMessage {
  final Uint8List sender;
  final int time;
  final Uint8List body;
  final Uint8List id;

  const OpenedMessage(this.sender, this.time, this.body, this.id);
}

/// Options for [seal]; [esk] and [pad] exist so tests can pin them, mirroring
/// the spec's fixed-ephemeral vectors.
class SealOptions {
  final Uint8List? esk;
  final bool pad;

  const SealOptions({this.esk, this.pad = true});
}

// §5.2 and §5.3. The ephemeral key is thrown away after sealing, so the sender
// cannot decrypt what they sent.
Uint8List seal(SmolIdentity identity, Uint8List recipient, Uint8List body,
    [int? when, SealOptions opts = const SealOptions()]) {
  final esk = opts.esk ?? randomBytes(keyLen);
  final epk = x25519Base(esk);
  final key = hkdfSha256(x25519(esk, ed25519ToX25519(recipient)),
      concat([epk, recipient]), _label.seal);
  final header = concat([
    Uint8List.fromList([1]),
    identity.publicKey,
    i64be(BigInt.from(when ?? nowSeconds())),
    u32be(body.length),
  ]);
  var plaintext = concat([
    header,
    body,
    ed25519Sign(identity.seed, concat([_label.msg, recipient, epk, header, body])),
  ]);
  if (opts.pad) {
    plaintext = concat(
        [plaintext, Uint8List((padTo - plaintext.length % padTo) % padTo)]);
  }
  final aad = concat([utf8Bytes("SMOL"), Uint8List.fromList([1]), recipient, epk]);
  return concat([aad, aeadEncrypt(key, Uint8List(12), plaintext, aad)]);
}

// Inverse of seal(); throws unless the signature and the recipient both check
// out. [identities] may include retired keys, per §7.
OpenedMessage unseal(List<SmolIdentity> identities, Uint8List envelope) {
  if (envelope.length < envelopeHeader + 16) {
    throw const SmolError("envelope too short");
  }
  final magic = utf8Bytes("SMOL");
  for (var i = 0; i < 4; i++) {
    if (magic[i] != envelope[i]) {
      throw const SmolError("not a Smol Mail envelope");
    }
  }
  if (envelope[4] != 1) {
    throw SmolError("unsupported envelope version ${envelope[4]}");
  }
  final to = envelope.sublist(5, 37), epk = envelope.sublist(37, 69);
  final sealed = envelope.sublist(69);
  SmolIdentity? me;
  for (final i in identities) {
    if (timingSafeEqual(i.publicKey, to)) {
      me = i;
      break;
    }
  }
  if (me == null) {
    throw SmolError(
        "addressed to ${b32encode(to).substring(0, 16)}…, not one of our keys");
  }
  final key = hkdfSha256(
      x25519(ed25519SeedToX25519(me.seed), epk), concat([epk, to]), _label.seal);
  Uint8List plaintext;
  try {
    plaintext = aeadDecrypt(key, Uint8List(12), sealed, envelope.sublist(0, envelopeHeader));
  } on SmolError {
    throw const SmolError("decryption failed: wrong key or corrupt envelope");
  }
  final r = Reader(plaintext);
  if (r.u8() != 1) throw const SmolError("unsupported payload version");
  final sender = r.take(keyLen);
  final when = r.i64();
  final bodyLen = r.u32();
  if (bodyLen > r.left) {
    throw const SmolError("payload body length exceeds the payload");
  }
  final body = r.take(bodyLen);
  final signature = r.take(sigLen); // trailing bytes are padding
  if (!ed25519Verify(sender,
      concat([_label.msg, to, epk, plaintext.sublist(0, payloadHeader), body]),
      signature)) {
    throw const SmolError("signature does not verify");
  }
  return OpenedMessage(sender, when, body, messageId(envelope));
}

// --- body frontmatter (§5.5) ---------------------------------------------------

final _fmKey = RegExp(r"^[A-Za-z0-9-]{1,64}$");

class Frontmatter {
  final Map<String, String> fields;
  final String body;

  const Frontmatter(this.fields, this.body);
}

// A flat `Key: value` block, deliberately not YAML. Any malformed line
// invalidates the whole block, which is then returned as ordinary body text:
// frontmatter fails closed toward display, never toward silent discard.
Frontmatter parseFrontmatter(String text) {
  if (!text.startsWith("---\n")) return Frontmatter(const {}, text);
  final lines = text.split("\n");
  final close = lines.indexOf("---", 1);
  if (close < 0) return Frontmatter(const {}, text);
  final block = lines.sublist(1, close);
  final rest = lines.sublist(close + 1).join("\n");
  var encoded = 0;
  for (final line in block) {
    encoded += utf8Bytes(line).length + 1;
  }
  if (block.length > _frontmatterKeys || encoded > _frontmatterMax) {
    return Frontmatter(const {}, text);
  }
  final fields = <String, String>{};
  for (final line in block) {
    final colon = line.indexOf(":");
    final head = colon < 0 ? "" : line.substring(0, colon);
    if (colon < 0 || !_fmKey.hasMatch(head)) {
      return Frontmatter(const {}, text);
    }
    // first occurrence wins
    fields.putIfAbsent(head, () => line.substring(colon + 1).trim());
  }
  return Frontmatter(fields, rest);
}

// Emit a block only when needed, including to escape a body that genuinely
// begins with `---` (§5.5).
String buildFrontmatter(Map<String, String> fields, String body) {
  final entries = fields.entries.where((e) => e.value.isNotEmpty).toList();
  if (entries.isEmpty && !body.startsWith("---\n")) return body;
  final block = entries.map((e) => "${e.key}: ${e.value}\n").join();
  return "---\n$block---\n$body";
}

// --- key rotation (§7) ---------------------------------------------------------

Uint8List makeCert(SmolIdentity oldIdentity, Uint8List newSeed, [int? when]) {
  final newPub = ed25519PublicKey(newSeed);
  final time = i64be(BigInt.from(when ?? nowSeconds()));
  return concat([
    oldIdentity.publicKey,
    newPub,
    time,
    ed25519Sign(oldIdentity.seed,
        concat([_label.rotate, oldIdentity.publicKey, newPub, time])),
  ]);
}

// Accept a key change only when a signed chain leads from the key we hold to
// the one the server now returns (§7).
bool walkChain(Uint8List pinned, Uint8List current, List<Uint8List> chain) {
  if (timingSafeEqual(pinned, current)) return true;
  if (chain.isEmpty || chain.length > maxChain) return false;
  var key = pinned;
  var started = false;
  for (final cert in chain) {
    final old = cert.sublist(0, 32), next = cert.sublist(32, 64);
    final when = cert.sublist(64, 72), sig = cert.sublist(72);
    if (!started) {
      if (!timingSafeEqual(old, key)) continue; // a link predating the key we hold
      started = true;
    } else if (!timingSafeEqual(old, key)) {
      return false; // the chain is not continuous
    }
    if (!ed25519Verify(old, concat([_label.rotate, old, next, when]), sig)) {
      return false;
    }
    key = next;
  }
  return started && timingSafeEqual(key, current);
}

// --- framing and operations (§4, §6) -------------------------------------------

/// An ordered byte pipe (TCP socket, or an in-memory queue in tests).
abstract class Wire {
  void send(Uint8List bytes);

  void close();

  Future<Uint8List> readExact(int n);
}

// One Noise session: application frames split across u16-prefixed Noise
// messages, requests and responses as in §6.1.
class Session {
  final Wire wire;
  final SessionCipher send, recv;

  Session(this.wire, this.send, this.recv);

  Future<Uint8List> _readNoise() async {
    final head = await wire.readExact(2);
    final length = (head[0] << 8) | head[1];
    if (length < 16) throw SmolError("server sent a $length-byte Noise message");
    return recv.decrypt(await wire.readExact(length));
  }

  Future<Response> call(int op, [Uint8List? body]) async {
    final payload = body ?? Uint8List(0);
    final frame = concat([u32be(1 + payload.length), Uint8List.fromList([op]), payload]);
    if (frame.length > maxFrame + 4) {
      throw const SmolError("request exceeds the maximum frame size");
    }
    for (var off = 0; off < frame.length; off += noisePayload) {
      final packet = send.encrypt(frame.sublist(off, min(off + noisePayload, frame.length)));
      wire.send(concat([u16be(packet.length), packet]));
    }
    var length = -1;
    var have = <int>[];
    while (length < 0 || have.length < 4 + length) {
      have.addAll(await _readNoise());
      if (length < 0 && have.length >= 4) {
        length = (have[0] << 24) | (have[1] << 16) | (have[2] << 8) | have[3];
        // §6.1: the shortest response is a type byte and a status byte.
        if (length < 2 || length > maxFrame) {
          throw SmolError("server sent a frame of length $length");
        }
      }
    }
    final payloadOut = Uint8List.fromList(have.sublist(4, 4 + length));
    // §6.1: a response reuses the request's type byte. A mismatch means the
    // session desynchronised, which must not be mistaken for a status.
    if (payloadOut[0] != op) {
      throw SmolError(
          "server answered op 0x${payloadOut[0].toRadixString(16)}, expected 0x${op.toRadixString(16)}");
    }
    return Response(payloadOut[1], payloadOut.sublist(2));
  }
}

class Response {
  final int status;
  final Uint8List body;

  const Response(this.status, this.body);
}

class OpenedSession {
  final Session session;
  final Uint8List serverStatic;
  final bool pinned;
  final Uint8List handshakeHash;

  const OpenedSession(this.session, this.serverStatic, this.pinned, this.handshakeHash);
}

// Handshake plus §4 pinning. Returns the session, the server's static key as
// revealed by the handshake, and whether that key was already pinned.
Future<OpenedSession> openSession(Wire wire, String host,
    {Uint8List? pinned}) async {
  final nx = NxInitiator();
  final m1 = nx.writeMessage1();
  wire.send(concat([u16be(m1.length), m1]));
  final head = await wire.readExact(2);
  final result = nx.readMessage2(await wire.readExact((head[0] << 8) | head[1]));
  if (pinned != null && !timingSafeEqual(pinned, result.serverStatic)) {
    throw SmolError("$host presented a different key than the one pinned\n"
        "  pinned:    ${b32encode(pinned)}\n"
        "  presented: ${b32encode(result.serverStatic)}");
  }
  return OpenedSession(
      Session(wire, result.send, result.recv),
      result.serverStatic,
      pinned != null,
      result.handshakeHash);
}

void expectOk(int status, String what) {
  if (status != 0) {
    throw SmolError("$what failed: ${statusName(status)} ($status)");
  }
}

// §4 session authentication: sign the handshake hash, which binds the
// signature to this session's server ephemeral and cannot be replayed.
Future<void> authenticate(
    Session session, Uint8List handshakeHash, String username, SmolIdentity identity) async {
  final name = utf8Bytes(username);
  if (name.length > 255) throw const SmolError("username too long");
  final body = concat([
    Uint8List.fromList([name.length]),
    name,
    identity.publicKey,
    ed25519Sign(identity.seed, concat([_label.auth, handshakeHash])),
  ]);
  expectOk((await session.call(opAuth, body)).status, "authentication");
}

class Resolved {
  final Uint8List identity;
  final List<Uint8List> chain;

  const Resolved(this.identity, this.chain);
}

// RESOLVE, returning the current key and its rotation chain (§6.1).
Future<Resolved> resolveOp(Session session, String user) async {
  final name = utf8Bytes(user);
  if (name.length > 255) throw const SmolError("username too long");
  final response =
      await session.call(opResolve, concat([Uint8List.fromList([name.length]), name]));
  expectOk(response.status, "resolving $user");
  final r = Reader(response.body);
  return Resolved(
      r.take(keyLen),
      List.generate(r.u8(), (_) => r.take(certLen)));
}

Future<Uint8List> sendOp(Session session, Uint8List envelope) async {
  final response = await session.call(opSend, envelope);
  expectOk(response.status, "sending");
  return response.body.length == idLen
      ? response.body
      : messageId(envelope);
}

class FetchedRecord {
  final Uint8List id;
  final int receivedAt;
  final Uint8List envelope;

  const FetchedRecord(this.id, this.receivedAt, this.envelope);
}

Future<List<FetchedRecord>> fetchOp(Session session) async {
  final response = await session.call(opFetch);
  expectOk(response.status, "fetching");
  final r = Reader(response.body);
  return List.generate(r.u16(), (_) {
    final id = r.take(idLen);
    final receivedAt = r.i64();
    return FetchedRecord(id, receivedAt, r.take(r.u32()));
  });
}

Future<int> deleteOp(Session session, List<Uint8List> ids) async {
  if (ids.length > 0xffff) throw const SmolError("too many ids for one DELETE");
  final body = concat([u16be(ids.length), ...ids]);
  final response = await session.call(opDelete, body);
  expectOk(response.status, "acknowledging");
  return Reader(response.body).u16();
}

class RegisterOptions {
  final String token;
  final Uint8List? cert;

  const RegisterOptions({this.token = "", this.cert});
}

Future<void> registerOp(
    Session session, String username, SmolIdentity identity,
    [RegisterOptions opts = const RegisterOptions()]) async {
  final name = utf8Bytes(username);
  final tokenBytes = utf8Bytes(opts.token);
  final cert = opts.cert ?? Uint8List(0);
  if (name.length > 255 || tokenBytes.length > 255 || cert.length > 255) {
    throw const SmolError("REGISTER field too long");
  }
  final body = concat([
    Uint8List.fromList([name.length]),
    name,
    identity.publicKey,
    Uint8List.fromList([tokenBytes.length]),
    tokenBytes,
    Uint8List.fromList([cert.length]),
    cert,
  ]);
  expectOk((await session.call(opRegister, body)).status, "registering $username");
}