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

import {
  aeadDecrypt, aeadEncrypt, concat, ed25519PublicKey, ed25519SeedToX25519,
  ed25519Sign, ed25519ToX25519, ed25519Verify, hkdfSha256, hmacSha256, randomBytes, sha256,
  timingSafeEqual, utf8Bytes, x25519, x25519Base,
} from "./crypto.js";
import { NxInitiator } from "./noise.js";

export const DEFAULT_PORT = 1961;
export const KEY_LEN = 32, SIG_LEN = 64, CERT_LEN = 200, ID_LEN = 32, TOKEN_LEN = 32;
export const MAX_FRAME = 1 << 20, NOISE_PAYLOAD = 65535 - 16, PAD_TO = 1024;
export const ENVELOPE_HEADER = 69, PAYLOAD_HEADER = 45, MAX_CHAIN = 16;
export const MAX_SKEW = 86400;  // §5.3: how far ahead of our clock a payload may be dated
export const FLAG_REQUESTS = 0x01;  // §6.1: set when a FETCH record missed an accept token
const FRONTMATTER_MAX = 4096, FRONTMATTER_KEYS = 64;

export const OP = { AUTH: 0x00, RESOLVE: 0x01, SEND: 0x02, FETCH: 0x03, DELETE: 0x04, REGISTER: 0x05 };
export const STATUS = { 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" };

export class SmolError extends Error {}

const LABEL = Object.freeze({
  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"), identity: utf8Bytes("smolmail/1 identity"),
  accept: utf8Bytes("smolmail/1 accept"), mac: utf8Bytes("smolmail/1 mac"),
  register: utf8Bytes("smolmail/1 register"),
});

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

const B32 = "ABCDEFGHIJKLMNOPQRSTUVWXYZ234567";

export function b32encode(bytes) {
  let out = "", value = 0, bits = 0;
  for (const b of bytes) {
    value = (value << 8) | b;
    bits += 8;
    while (bits >= 5) { bits -= 5; out += B32[(value >>> bits) & 31]; }
  }
  if (bits) out += B32[(value << (5 - bits)) & 31];
  return out.toLowerCase();
}

export function b32decode(text) {
  let out = [], value = 0, bits = 0;
  for (const ch of text.trim().toUpperCase().replace(/=+$/, "")) {
    const idx = B32.indexOf(ch);
    if (idx < 0) throw new SmolError(`invalid base32 character '${ch}'`);
    value = (value << 5) | idx;
    bits += 5;
    if (bits >= 8) { bits -= 8; out.push((value >>> bits) & 0xff); }
  }
  return Uint8Array.from(out);
}

// §3: the first 20 base32 characters of the identity, in groups of four.
export function fingerprint(identity) {
  const s = b32encode(identity).slice(0, 20);
  return s.match(/.{4}/g).join(" ");
}

export const u16BE = (n) => Uint8Array.of((n >> 8) & 0xff, n & 0xff);

export const u32BE = (n) => Uint8Array.of((n >>> 24) & 0xff, (n >>> 16) & 0xff, (n >>> 8) & 0xff, n & 0xff);

export const i64BE = (n) => {
  const out = new Uint8Array(8);
  new DataView(out.buffer).setBigInt64(0, BigInt(n), false);
  return out;
};

export const nowSeconds = () => Math.floor(Date.now() / 1000);

// Fail-closed reader; every parse raises rather than reading past the end.
export class Reader {
  constructor(buf) { this.buf = buf; this.pos = 0; }
  take(n) {
    if (n < 0 || this.pos + n > this.buf.length) throw new SmolError("truncated message");
    this.pos += n;
    return this.buf.slice(this.pos - n, this.pos);
  }
  u8() { return this.take(1)[0]; }
  u16() { const b = this.take(2); return (b[0] << 8) | b[1]; }
  u32() { return new DataView(this.take(4).buffer).getUint32(0); }
  i64() { return new DataView(this.take(8).buffer).getBigInt64(0); }
  get left() { return this.buf.length - this.pos; }
  get done() { return this.pos === this.buf.length; }
}

// --- identity (§2) ------------------------------------------------------------

// An Ed25519 keypair with the X25519 agreement keys derived from it.
export function identityFromSeed(seed) {
  if (seed.length !== KEY_LEN) throw new SmolError(`identity seed must be ${KEY_LEN} bytes`);
  return { seed, publicKey: ed25519PublicKey(seed) };
}

export function newMaster() {
  return randomBytes(KEY_LEN);
}

// §2: the only secret a user holds. Every rotation index's signing seed, and
// the accept key, are derived from it with HKDF.
export function identitySeed(master, index) {
  return hkdfSha256(master, new Uint8Array(0), concat(LABEL.identity, u32BE(index)));
}

export function acceptKeyFor(master) {
  return hkdfSha256(master, new Uint8Array(0), LABEL.accept);
}

// §5.8: the token this account issues to one correspondent, independent of
// the rotation index so it survives the owner's key rotation.
export function tokenFor(master, correspondentIdentity) {
  return hmacSha256(acceptKeyFor(master), correspondentIdentity);
}

// §5.8: what a sender attaches to SEND to reach the recipient's main tier.
export function acceptMac(token, id) {
  return hmacSha256(token, concat(LABEL.mac, id));
}

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

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

export function parseAddress(text) {
  text = text.trim();
  let identity = null;
  if (text.startsWith("smol://")) {
    const rest = text.slice("smol://".length);
    const slash = rest.lastIndexOf("/");
    if (slash < 0) throw new SmolError(`${text}: smol:// address carries no key`);
    identity = b32decode(rest.slice(slash + 1));
    if (identity.length !== KEY_LEN)
      throw new SmolError(`${text}: key is ${identity.length} bytes, expected ${KEY_LEN}`);
    text = rest.slice(0, slash);
  }
  const m = ADDRESS.exec(text.toLowerCase());
  if (!m) throw new SmolError(`'${text}' is not a valid address`);
  const { user, host } = m.groups;
  if ("._-".includes(user[0]) || "._-".includes(user.at(-1)))
    throw new SmolError(`${user} may not begin or end with a separator`);
  const port = m.groups.port ? Number(m.groups.port) : DEFAULT_PORT;
  return {
    user, host, port, identity,
    get short() { return `${user}@${host}${port === DEFAULT_PORT ? "" : ":" + port}`; },
    uri: (key) => `smol://${user}@${host}${port === DEFAULT_PORT ? "" : ":" + port}/${b32encode(key)}`,
  };
}

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

// §5.4: derived from the envelope so no sender can choose it; used whole,
// nothing truncates it.
export function messageId(envelope) {
  return sha256(concat(LABEL.id, envelope));
}

// §5.2 and §5.3. The ephemeral key is thrown away after sealing, so the sender
// cannot decrypt what they sent; opts.esk exists only so tests can pin it.
export function seal(identity, recipient, body, when = nowSeconds(), opts = {}) {
  const esk = opts.esk ?? randomBytes(KEY_LEN);
  const epk = x25519Base(esk);
  const key = hkdfSha256(x25519(esk, ed25519ToX25519(recipient)), concat(epk, recipient), LABEL.seal);
  const header = concat(Uint8Array.of(1), identity.publicKey, i64BE(when), u32BE(body.length));
  let plaintext = concat(header, body,
    ed25519Sign(identity.seed, concat(LABEL.msg, recipient, epk, header, body)));
  if (opts.pad !== false)
    plaintext = concat(plaintext, new Uint8Array((PAD_TO - plaintext.length % PAD_TO) % PAD_TO));
  const aad = concat(utf8Bytes("SMOL"), Uint8Array.of(1), recipient, epk);
  return concat(aad, aeadEncrypt(key, new Uint8Array(12), plaintext, aad));
}

// Inverse of seal(); throws unless the signature and the recipient both check
// out. `identities` may include retired keys, per §7.
export function unseal(identities, envelope) {
  if (envelope.length < ENVELOPE_HEADER + 16) throw new SmolError("envelope too short");
  if (utf8Bytes("SMOL").some((b, i) => b !== envelope[i])) throw new SmolError("not a Smol Mail envelope");
  if (envelope[4] !== 1) throw new SmolError(`unsupported envelope version ${envelope[4]}`);
  const to = envelope.slice(5, 37), epk = envelope.slice(37, 69), sealed = envelope.slice(69);
  const me = identities.find(i => timingSafeEqual(i.publicKey, to));
  if (!me) throw new SmolError(`addressed to ${b32encode(to).slice(0, 16)}…, not one of our keys`);
  const key = hkdfSha256(x25519(ed25519SeedToX25519(me.seed), epk), concat(epk, to), LABEL.seal);
  let plaintext;
  try {
    plaintext = aeadDecrypt(key, new Uint8Array(12), sealed, envelope.slice(0, ENVELOPE_HEADER));
  } catch {
    throw new SmolError("decryption failed: wrong key or corrupt envelope");
  }
  const r = new Reader(plaintext);
  if (r.u8() !== 1) throw new SmolError("unsupported payload version");
  const sender = r.take(KEY_LEN), when = r.i64(), bodyLen = r.u32();
  if (bodyLen > r.left) throw new SmolError("payload body length exceeds the payload");
  const body = r.take(bodyLen), signature = r.take(SIG_LEN);  // trailing bytes are padding
  if (!ed25519Verify(sender, concat(LABEL.msg, to, epk, plaintext.slice(0, PAYLOAD_HEADER), body), signature))
    throw new SmolError("signature does not verify");
  if (when > BigInt(nowSeconds() + MAX_SKEW)) throw new SmolError("payload is dated in the future");
  return { sender, time: Number(when), body, id: messageId(envelope) };
}

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

const FM_KEY = /^[A-Za-z0-9-]{1,64}$/;

// 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. Keys
// are compared case-insensitively (§5.5), so they are returned lowercased.
export function parseFrontmatter(text) {
  if (!text.startsWith("---\n")) return { fields: {}, body: text };
  const lines = text.split("\n");
  const close = lines.indexOf("---", 1);
  if (close < 0) return { fields: {}, body: text };
  const block = lines.slice(1, close), rest = lines.slice(close + 1).join("\n");
  const encoded = block.map(line => utf8Bytes(line).length + 1);
  if (block.length > FRONTMATTER_KEYS || encoded.reduce((a, b) => a + b, 0) > FRONTMATTER_MAX)
    return { fields: {}, body: text };
  const fields = {};
  for (let i = 0; i < block.length; i++) {
    const line = block[i], colon = line.indexOf(":");
    if (colon < 0 || !FM_KEY.test(line.slice(0, colon))) return { fields: {}, body: text };
    const key = line.slice(0, colon).toLowerCase();
    if (!(key in fields)) fields[key] = line.slice(colon + 1).trim();  // first occurrence wins
  }
  return { fields, body: rest };
}

// Emit a block only when needed, including to escape a body that genuinely
// begins with `---` (§5.5).
export function buildFrontmatter(fields, body) {
  if (!fields.length && !body.startsWith("---\n")) return body;
  return `---\n${fields.map(([k, v]) => `${k}: ${v}\n`).join("")}---\n${body}`;
}

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

// old_pub 32 || new_pub 32 || time 8 || sig_old 64 || sig_new 64. Both keys
// sign, so the old key alone cannot hand the username to a key nobody
// controls; the username is covered but not carried, so a verifier always
// supplies the one it is checking.
export function makeCert(username, oldIdentity, newSeed, when = nowSeconds()) {
  const newPub = ed25519PublicKey(newSeed), time = i64BE(when);
  const signed = concat(LABEL.rotate, utf8Bytes(username), oldIdentity.publicKey, newPub, time);
  return concat(oldIdentity.publicKey, newPub, time,
    ed25519Sign(oldIdentity.seed, signed), ed25519Sign(newSeed, signed));
}

// Accept a key change only when a signed chain leads from the key we hold to
// the one the server now returns, both keys signing each link (§7).
export function walkChain(username, pinned, current, chain) {
  const same = (a, b) => timingSafeEqual(a, b);
  if (same(pinned, current)) return true;
  if (!chain.length || chain.length > MAX_CHAIN) return false;
  let key = pinned, started = false;
  for (const cert of chain) {
    const old = cert.slice(0, 32), next = cert.slice(32, 64), when = cert.slice(64, 72);
    const sigOld = cert.slice(72, 136), sigNew = cert.slice(136, 200);
    if (!started) {
      if (!same(old, key)) continue;  // a link predating the key we hold
      started = true;
    } else if (!same(old, key)) {
      return false;  // the chain is not continuous
    }
    const signed = concat(LABEL.rotate, utf8Bytes(username), old, next, when);
    if (!ed25519Verify(old, signed, sigOld) || !ed25519Verify(next, signed, sigNew)) return false;
    key = next;
  }
  return started && same(key, current);
}

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

// Reassembles a byte stream (WebSocket or TCP) into exact-length reads.
class ByteStream {
  constructor(stream, timeoutMs = 0) {
    this.chunks = [];
    this.length = 0;
    this.closed = null;
    this.waiters = [];
    this.timeoutMs = timeoutMs;
    stream.onData = (data) => { this.chunks.push(data); this.length += data.length; this.wake(); };
    stream.onClose = () => this.fail(new SmolError("server closed the connection"));
  }

  wake() {
    this.waiters = this.waiters.filter(w => {
      if (this.closed) { w.reject(this.closed); return false; }
      if (this.length >= w.need) { w.resolve(); return false; }
      return true;
    });
  }

  fail(error) {
    this.closed = error;
    this.waiters.forEach(w => w.reject(error));
    this.waiters = [];
  }

  async readExact(n) {
    if (this.closed) throw this.closed;
    if (this.length < n) {
      await new Promise((resolve, reject) => {
        const waiter = { need: n, resolve, reject };
        this.waiters.push(waiter);
        if (!this.timeoutMs) return;
        // Settings' timeout, applied per read: a stalled handshake or request
        // otherwise waits here forever, since nothing else ever rejects it.
        setTimeout(() => {
          if (!this.waiters.includes(waiter)) return;  // already settled elsewhere
          this.waiters = this.waiters.filter(w => w !== waiter);
          reject(new SmolError(`no response from the server within ${this.timeoutMs / 1000}s`));
        }, this.timeoutMs);
      });
      if (this.closed) throw this.closed;
    }
    const out = new Uint8Array(n);
    let off = 0;
    while (off < n) {
      const chunk = this.chunks[0];
      const take = Math.min(chunk.length, n - off);
      out.set(off ? chunk.slice(0, take) : chunk.subarray(0, take), off);
      if (take === chunk.length) this.chunks.shift();
      else this.chunks[0] = chunk.subarray(take);
      off += take;
      this.length -= take;
    }
    return out;
  }
}

// One Noise session: application frames split across u16-prefixed Noise
// messages, requests and responses as in §6.1.
export class Session {
  constructor(wire, stream, send, recv) {
    this.wire = wire;
    this.stream = stream;
    this.send = send;
    this.recv = recv;
  }

  async readNoise() {
    const head = await this.wire.readExact(2);
    const length = (head[0] << 8) | head[1];
    if (length < 16) throw new SmolError(`server sent a ${length}-byte Noise message`);
    return this.recv.decrypt(await this.wire.readExact(length));
  }

  async call(op, body = new Uint8Array(0)) {
    const frame = concat(u32BE(1 + body.length), Uint8Array.of(op), body);
    if (frame.length > MAX_FRAME + 4) throw new SmolError("request exceeds the maximum frame size");
    for (let off = 0; off < frame.length; off += NOISE_PAYLOAD) {
      const packet = this.send.encrypt(frame.slice(off, off + NOISE_PAYLOAD));
      this.stream.send(concat(u16BE(packet.length), packet));
    }
    let length = -1;
    let have = new Uint8Array(0);
    while (length < 0 || have.length < 4 + length) {
      have = concat(have, await this.readNoise());
      if (length < 0 && have.length >= 4) {
        length = new DataView(have.buffer).getUint32(0);
        // §6.1: the shortest response is a type byte and a status byte.
        if (length < 2 || length > MAX_FRAME)
          throw new SmolError(`server sent a frame of length ${length}`);
      }
    }
    const payload = have.slice(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 (payload[0] !== op)
      throw new SmolError(`server answered op 0x${payload[0].toString(16)}, expected 0x${op.toString(16)}`);
    return { status: payload[1], body: payload.slice(2) };  // op echo, status, body (§6.1)
  }
}

// Handshake plus §4 pinning. Returns the session, the server's static key as
// revealed by the handshake, and whether that key was already pinned.
export async function openSession(stream, host, pinned = null, timeoutMs = 0) {
  const wire = new ByteStream(stream, timeoutMs);
  const nx = new NxInitiator();
  const m1 = nx.writeMessage1();
  stream.send(concat(u16BE(m1.length), m1));
  const head = await wire.readExact(2);
  const ciphers = nx.readMessage2(await wire.readExact((head[0] << 8) | head[1]));
  if (pinned && !timingSafeEqual(pinned, nx.serverStatic))
    throw new SmolError(`${host} presented a different key than the one pinned\n` +
      `  pinned:    ${b32encode(pinned)}\n  presented: ${b32encode(nx.serverStatic)}`);
  return {
    session: new Session(wire, stream, ciphers.send, ciphers.recv),
    serverStatic: nx.serverStatic,
    pinned: pinned !== null,
    handshakeHash: nx.handshakeHash,
  };
}

export function expectOk(status, what) {
  if (status !== 0) throw new SmolError(`${what} failed: ${STATUS[status] ?? status} (${status})`);
}

// §4 session authentication: sign the handshake hash, which binds the
// signature to this session's server ephemeral and cannot be replayed, and
// push the accept token set (§5.8). `sync = 0` leaves the server's stored set
// untouched and `tokens` MUST then be empty; `sync = 1` replaces it exactly.
// Returns the number of accept tokens the server now holds.
export async function authenticate(session, handshakeHash, username, identity, { sync = 0, tokens = [] } = {}) {
  const name = utf8Bytes(username);
  if (name.length > 255) throw new SmolError("username too long");
  if (tokens.length > 0xffff) throw new SmolError("too many accept tokens for one AUTH");
  const body = concat(Uint8Array.of(name.length), name, identity.publicKey,
    ed25519Sign(identity.seed, concat(LABEL.auth, handshakeHash)),
    Uint8Array.of(sync), u16BE(tokens.length), ...tokens);
  const { status, body: reply } = await session.call(OP.AUTH, body);
  expectOk(status, "authentication");
  return reply.length >= 2 ? new Reader(reply).u16() : 0;
}

// RESOLVE, returning the current key and its rotation chain (§6.1).
export async function resolveOp(session, user) {
  const name = utf8Bytes(user);
  if (name.length > 255) throw new SmolError("username too long");
  const { status, body } = await session.call(OP.RESOLVE, concat(Uint8Array.of(name.length), name));
  expectOk(status, `resolving ${user}`);
  const r = new Reader(body);
  return { identity: r.take(KEY_LEN), chain: Array.from({ length: r.u8() }, () => r.take(CERT_LEN)) };
}

// §5.8: `mac` is the sender's proof of an accept token, absent or TOKEN_LEN bytes.
export async function sendOp(session, envelope, mac = null) {
  const macBytes = mac ?? new Uint8Array(0);
  if (macBytes.length && macBytes.length !== TOKEN_LEN)
    throw new SmolError(`accept MAC must be ${TOKEN_LEN} bytes`);
  const body = concat(Uint8Array.of(macBytes.length), macBytes, envelope);
  const { status, body: reply } = await session.call(OP.SEND, body);
  expectOk(status, "sending");
  return reply.length === ID_LEN ? reply : messageId(envelope);
}

// §6.1: pages forward from a cursor; an all-zero id starts at the beginning.
export async function fetchOp(session, afterReceivedAt = 0, afterId = new Uint8Array(ID_LEN)) {
  const body = concat(i64BE(afterReceivedAt), afterId);
  const { status, body: reply } = await session.call(OP.FETCH, body);
  expectOk(status, "fetching");
  const r = new Reader(reply);
  return Array.from({ length: r.u16() }, () => {
    const id = r.take(ID_LEN), receivedAt = Number(r.i64()), flags = r.u8();
    return { id, receivedAt, flags, envelope: r.take(r.u32()), isRequest: Boolean(flags & FLAG_REQUESTS) };
  });
}

export async function deleteOp(session, ids) {
  if (ids.length > 0xffff) throw new SmolError("too many ids for one DELETE");
  const body = concat(u16BE(ids.length), ...ids);
  const { status, body: reply } = await session.call(OP.DELETE, body);
  expectOk(status, "acknowledging");
  return new Reader(reply).u16();
}

// §6.1: the signature is proof of possession, bound to the server that will
// store the binding so it cannot be replayed to another server.
export function registerSigned(serverStatic, username, identity) {
  return concat(LABEL.register, serverStatic, utf8Bytes(username), identity);
}

export async function registerOp(session, serverStatic, { username, identity, token = "", cert = null }) {
  const name = utf8Bytes(username);
  const tokenBytes = utf8Bytes(token);
  if (name.length > 255 || tokenBytes.length > 255 || cert && cert.length > 255)
    throw new SmolError("REGISTER field too long");
  const body = concat(Uint8Array.of(name.length), name, identity.publicKey,
    ed25519Sign(identity.seed, registerSigned(serverStatic, username, identity.publicKey)),
    Uint8Array.of(tokenBytes.length), tokenBytes,
    Uint8Array.of(cert ? cert.length : 0), cert ?? new Uint8Array(0));
  const { status } = await session.call(OP.REGISTER, body);
  expectOk(status, `registering ${username}`);
}