// 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 {
  aeadDecrypt, aeadEncrypt, concat, ed25519PublicKey, ed25519SeedToX25519,
  ed25519Sign, ed25519ToX25519, ed25519Verify, hkdfSha256, 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 = 136, ID_LEN = 16;
export const MAX_FRAME = 1 << 20, NOISE_PAYLOAD = 65535 - 16, PAD_TO = 1024;
export const ENVELOPE_HEADER = 69, PAYLOAD_HEADER = 45, MAX_CHAIN = 16;
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"),
});

// --- 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 newIdentity() {
  return identityFromSeed(randomBytes(KEY_LEN));
}

// --- 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) ----------------------------------------------------

export function messageId(envelope) {
  return sha256(concat(LABEL.id, envelope)).slice(0, ID_LEN);
}

// §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");
  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.
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);
    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) -------------------------------------------------------

export function makeCert(oldIdentity, newSeed, when = nowSeconds()) {
  const newPub = ed25519PublicKey(newSeed), time = i64BE(when);
  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).
export function walkChain(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), sig = cert.slice(72);
    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
    }
    if (!ed25519Verify(old, concat(LABEL.rotate, old, next, when), sig)) 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) {
    this.chunks = [];
    this.length = 0;
    this.closed = null;
    this.waiters = [];
    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) => this.waiters.push({ need: n, resolve, reject }));
      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) {
  const wire = new ByteStream(stream);
  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.
export async function authenticate(session, handshakeHash, username, identity) {
  const name = utf8Bytes(username);
  if (name.length > 255) throw new SmolError("username too long");
  const body = concat(Uint8Array.of(name.length), name, identity.publicKey,
    ed25519Sign(identity.seed, concat(LABEL.auth, handshakeHash)));
  expectOk((await session.call(OP.AUTH, body)).status, "authentication");
}

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

export async function sendOp(session, envelope) {
  const { status, body } = await session.call(OP.SEND, envelope);
  expectOk(status, "sending");
  return body.length === ID_LEN ? body : messageId(envelope);
}

export async function fetchOp(session) {
  const { status, body } = await session.call(OP.FETCH);
  expectOk(status, "fetching");
  const r = new Reader(body);
  return Array.from({ length: r.u16() }, () => {
    const id = r.take(ID_LEN), receivedAt = Number(r.i64());
    return { id, receivedAt, envelope: r.take(r.u32()) };
  });
}

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

export async function registerOp(session, { 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,
    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}`);
}