// Headless driver for the end-to-end test: the same web/js modules the browser
// uses, over a TCP socket (or, with GSMOL_TRANSPORT=ws, through the bridge's
// WebSocket relay — the exact byte path the browser takes).
//
//   node test/client.mjs --state bob.json <command> [args]
//
// Commands mirror the reference smolmail.py: keygen, trust, register, resolve,
// import, accept, block, send, fetch, list, read, rotate. State lives in a
// JSON file; mail is stored sealed, like the browser store does.

import { existsSync, readFileSync, writeFileSync } from "node:fs";
import net from "node:net";
import { hex, timingSafeEqual, unhex, utf8Bytes } from "../web/js/crypto.js";
import * as proto from "../web/js/proto.js";

const { b32decode, b32encode, fingerprint } = proto;

const args = process.argv.slice(2);
const stateIndex = args.indexOf("--state");
const statePath = stateIndex >= 0 ? args[stateIndex + 1] : "client-state.json";
const command = args.find((a, i) => i !== stateIndex && i !== stateIndex + 1 && !a.startsWith("--"));
const rest = args.slice(args.indexOf(command) + 1);

const load = () => existsSync(statePath) ? JSON.parse(readFileSync(statePath, "utf8")) : {};
const save = (state) => writeFileSync(statePath, JSON.stringify(state));

function parseOption(name) {
  const i = rest.indexOf(name);
  return i >= 0 ? rest[i + 1] : null;
}

// --- transports -------------------------------------------------------------

function connectTcp(host, port) {
  return new Promise((resolve, reject) => {
    const socket = net.connect({ host, port });
    socket.on("error", reject);
    socket.on("connect", () => {
      const stream = {
        send: (bytes) => socket.write(bytes),
        close: () => socket.end(),
        onData: null,
        onClose: null,
      };
      socket.on("data", (chunk) => stream.onData?.(new Uint8Array(chunk)));
      socket.on("close", () => stream.onClose?.());
      resolve(stream);
    });
  });
}

function connectWs(host, port) {
  const bridge = process.env.GSMOL_BRIDGE || "ws://127.0.0.1:8096";
  return new Promise((resolve, reject) => {
    const socket = new WebSocket(`${bridge}/tcp/${host}/${port}`);
    socket.binaryType = "arraybuffer";
    const stream = {
      send: (bytes) => socket.send(bytes),
      close: () => socket.close(),
      onData: null,
      onClose: null,
    };
    socket.onerror = () => reject(new Error(`cannot reach the bridge at ${bridge}`));
    socket.onopen = () => resolve(stream);
    socket.onmessage = (event) => stream.onData?.(new Uint8Array(event.data));
    socket.onclose = () => stream.onClose?.();
  });
}

const connectStream = process.env.GSMOL_TRANSPORT === "ws" ? connectWs : connectTcp;

// --- session ------------------------------------------------------------------

async function open(host, port, { requirePin }) {
  const state = load();
  const pinned = state.servers?.[host] ? b32decode(state.servers[host]) : null;
  if (requirePin && !pinned)
    throw new proto.SmolError(`no pinned key for ${host}`);
  const stream = await connectStream(host, port);
  const opened = await proto.openSession(stream, host, pinned);
  if (!pinned)
    console.error(`warning: ${host} is not pinned; its key is ${b32encode(opened.serverStatic)}`);
  return opened;
}

// §2: the identity at the state's current rotation index, plus every earlier
// one — mail sealed to a superseded key is readable with nothing else.
function currentIdentity(state) {
  return proto.identityFromSeed(proto.identitySeed(unhex(state.master), state.rotations ?? 0));
}

function identities(state) {
  const master = unhex(state.master);
  const out = [];
  for (let n = state.rotations ?? 0; n >= 0; n--) out.push(proto.identityFromSeed(proto.identitySeed(master, n)));
  return out;
}

function cursor(state) {
  return [state.afterTime ?? 0, state.afterId ? unhex(state.afterId) : new Uint8Array(proto.ID_LEN)];
}

// §4/§5.8: the accept tokens to push with AUTH, and whether to push at all —
// a client that cannot vouch for its own set (freshly restored) must not
// replace the server's with an incomplete one.
function tokenSet(state) {
  if (state.syncOk === false) return { sync: 0, tokens: [] };
  const accepted = Object.entries(state.accepted || {}).filter(([, a]) => a.active)
    .sort(([a], [b]) => a.localeCompare(b));
  const master = unhex(state.master);
  return { sync: 1, tokens: accepted.map(([, a]) => proto.tokenFor(master, b32decode(a.identity))) };
}

function addressOfSigner(state, sender, replyToField) {
  const known = Object.entries(state.contacts || {}).find(([, c]) => c.key === b32encode(sender))?.[0];
  if (known) return known;
  if (!replyToField) return null;
  try {
    const parsed = proto.parseAddress(replyToField);
    if (parsed.identity && timingSafeEqual(parsed.identity, sender)) return parsed.short;
  } catch { /* malformed claim */ }
  return null;
}

// --- commands -------------------------------------------------------------------

const commands = {
  keygen() {
    const state = load();
    state.master = hex(cryptoRandom(32));
    state.rotations = 0;
    state.syncOk = true;
    save(state);
    const me = currentIdentity(state);
    console.log(`public key:  ${b32encode(me.publicKey)}`);
    console.log(`fingerprint: ${fingerprint(me.publicKey)}`);
  },

  trust() {
    const host = rest[0], key = b32decode(rest[1]);
    if (key.length !== 32) throw new proto.SmolError(`server key is ${key.length} bytes, expected 32`);
    const state = load();
    state.servers = { ...(state.servers || {}), [host]: b32encode(key) };
    save(state);
    console.log(`pinned ${host}  ${b32encode(key)}`);
  },

  async register() {
    const state = load();
    const me = currentIdentity(state);
    const addr = proto.parseAddress(rest[0]);
    const opened = await open(addr.host, addr.port, { requirePin: true });
    try {
      await proto.registerOp(opened.session, opened.serverStatic, {
        username: addr.user, identity: me, token: parseOption("--token") || "",
      });
    } finally {
      opened.session.stream.close();
    }
    state.account = { user: addr.user, host: addr.host, port: addr.port };
    save(state);
    console.log(`registered ${addr.short}`);
    console.log(`share: ${addr.uri(me.publicKey)}`);
  },

  async resolve() {
    const state = load();
    const addr = proto.parseAddress(rest[0]);
    const opened = await open(addr.host, addr.port, { requirePin: false });
    let current, chain;
    try {
      ({ identity: current, chain } = await proto.resolveOp(opened.session, addr.user));
    } finally {
      opened.session.stream.close();
    }
    const known = state.contacts?.[addr.short];
    if (!known) {
      state.contacts = { ...(state.contacts || {}), [addr.short]: { key: b32encode(current), verified: false } };
      save(state);
      console.log(`${addr.short}  ${b32encode(current)}`);
      console.log(`  pinned (trust on first use${opened.pinned ? "" : ", UNVERIFIED server"})`);
      return;
    }
    if (timingSafeEqual(b32decode(known.key), current)) return console.log(`${addr.short}: unchanged`);
    if (proto.walkChain(addr.user, b32decode(known.key), current, chain)) {
      state.contacts[addr.short] = { key: b32encode(current), verified: known.verified };
      save(state);
      console.log(`warning: ${addr.short} rotated its key; a signed chain confirms it`);
      return;
    }
    throw new proto.SmolError(`${addr.short} presents a different key with no valid rotation chain`);
  },

  import() {
    const state = load();
    const addr = proto.parseAddress(rest[0]);
    if (!addr.identity) throw new proto.SmolError("import needs a smol:// address carrying a key");
    state.contacts = {
      ...(state.contacts || {}),
      [addr.short]: { key: b32encode(addr.identity), verified: true },
    };
    save(state);
    console.log(`imported ${addr.short}  ${b32encode(addr.identity)}  (verified)`);
  },

  contacts() {
    const state = load();
    for (const [address, c] of Object.entries(state.contacts || {})) {
      const accepted = state.accepted?.[address];
      const tag = accepted ? (accepted.active ? "accepted" : "blocked") : "";
      console.log(`${address}  ${c.key}  ${c.verified ? "verified" : "tofu"}  ${tag}`.trimEnd());
    }
  },

  // §5.8: admit a contact to the main tier; their token travels in our next
  // message to them. Pushed to the server right away.
  async accept() {
    const state = load();
    const address = rest[0];
    const contact = state.contacts?.[address];
    if (!contact) throw new proto.SmolError(`no key for ${address} yet`);
    state.accepted = { ...(state.accepted || {}),
      [address]: { identity: state.accepted?.[address]?.identity ?? contact.key, active: true } };
    state.syncOk = true;
    save(state);
    const held = await pushTokens(state);
    console.log(`accepted ${address}; server now holds ${held} accept token(s)`);
  },

  // §5.8: withdraw a contact's accept token; their mail lands in requests
  // from their next message on.
  async block() {
    const state = load();
    const address = rest[0];
    if (!state.accepted?.[address]) throw new proto.SmolError(`${address} was never accepted`);
    state.accepted[address] = { ...state.accepted[address], active: false };
    save(state);
    const held = await pushTokens(state);
    console.log(`blocked ${address}; server now holds ${held} accept token(s)`);
  },

  async send() {
    const state = load();
    const me = currentIdentity(state);
    const master = unhex(state.master);
    const addr = proto.parseAddress(rest[0]);
    let recipient;
    if (addr.identity) {
      recipient = addr.identity;
      state.contacts = {
        ...(state.contacts || {}),
        [addr.short]: { key: b32encode(recipient), verified: true },
      };
    } else if (state.contacts?.[addr.short]) {
      recipient = b32decode(state.contacts[addr.short].key);
    } else {
      const opened = await open(addr.host, addr.port, { requirePin: false });
      try {
        ({ identity: recipient } = await proto.resolveOp(opened.session, addr.user));
      } finally {
        opened.session.stream.close();
      }
      state.contacts = { ...(state.contacts || {}), [addr.short]: { key: b32encode(recipient), verified: false } };
      console.error(`warning: ${addr.short} pinned (trust on first use)`);
    }
    save(state);

    const fields = [["Subject", parseOption("--subject")], ["In-Reply-To", parseOption("--reply-to")]]
      .filter(([, v]) => v);
    if (state.account && !rest.includes("--anonymous"))
      fields.push(["Reply-To", proto.parseAddress(`${state.account.user}@${state.account.host}` +
        (state.account.port === proto.DEFAULT_PORT ? "" : `:${state.account.port}`)).uri(me.publicKey)]);
    // §5.8: hand an accepted correspondent the token for our own mailbox.
    const accepted = state.accepted?.[addr.short];
    if (accepted?.active) fields.push(["Accept", b32encode(proto.tokenFor(master, b32decode(accepted.identity)))]);
    const body = utf8Bytes(proto.buildFrontmatter(fields, (parseOption("--body") || "") + "\n"));
    const envelope = proto.seal(me, recipient, body);
    // §5.8: our token for their mailbox, if they have given us one.
    const held = state.tokens?.[addr.short];
    const mac = held ? proto.acceptMac(b32decode(held.token), proto.messageId(envelope)) : null;
    const opened = await open(addr.host, addr.port, { requirePin: false });
    try {
      await proto.sendOp(opened.session, envelope, mac);
    } finally {
      opened.session.stream.close();
    }
    state.sent = [...(state.sent || []), {
      id: hex(proto.messageId(envelope)), recipient: addr.short,
      envelope: hex(proto.seal(me, me.publicKey, body)), sentAt: proto.nowSeconds(),
    }];
    save(state);
    console.log(`sent ${hex(proto.messageId(envelope)).slice(0, 16)} to ${addr.short}` +
      (mac ? " (accepted)" : ""));
  },

  async fetch() {
    const state = load();
    const me = currentIdentity(state);
    const account = proto.parseAddress(`${state.account.user}@${state.account.host}` +
      (state.account.port === proto.DEFAULT_PORT ? "" : `:${state.account.port}`));
    const ids = identities(state);
    const keep = rest.includes("--keep");
    if (rest.includes("--reset")) { state.afterTime = 0; state.afterId = hex(new Uint8Array(proto.ID_LEN)); }
    let [afterTime, afterId] = keep ? cursor(state) : [0, new Uint8Array(proto.ID_LEN)];
    const opened = await open(account.host, account.port, { requirePin: true });
    let total = 0, stored = 0, rejected = 0;
    try {
      const { sync, tokens } = tokenSet(state);
      await proto.authenticate(opened.session, opened.handshakeHash, account.user, me, { sync, tokens });
      for (;;) {
        const records = await proto.fetchOp(opened.session, afterTime, afterId);
        if (!records.length) break;
        const acked = [];
        for (const record of records) {
          total++;
          afterTime = record.receivedAt;
          afterId = record.id;
          let unsealed;
          try {
            if (!timingSafeEqual(proto.messageId(record.envelope), record.id))
              throw new proto.SmolError("id does not match the envelope");
            unsealed = proto.unseal(ids, record.envelope);
          } catch (err) {
            rejected++;
            console.error(`warning: ${hex(record.id)}: ${err.message}; left on server`);
            continue;
          }
          if (!state.inbox) state.inbox = [];
          if (!state.inbox.some(m => m.id === hex(record.id))) {
            state.inbox.push({
              id: hex(record.id), envelope: hex(record.envelope), receivedAt: record.receivedAt,
              isRequest: record.isRequest, keptOnServer: keep,
            });
            stored++;
            const { fields } = proto.parseFrontmatter(new TextDecoder().decode(unsealed.body));
            const raw = fields.accept;
            if (raw) {
              try {
                const token = b32decode(raw);
                if (token.length === proto.TOKEN_LEN) {
                  const address = addressOfSigner(state, unsealed.sender, fields["reply-to"]);
                  if (address) state.tokens = { ...(state.tokens || {}), [address]: { token: b32encode(token) } };
                }
              } catch { /* malformed token: ignore */ }
            }
          }
          acked.push(record.id);
        }
        if (keep) {
          state.afterTime = afterTime;
          state.afterId = hex(afterId);
        } else if (acked.length) {
          await proto.deleteOp(opened.session, acked);
        }
      }
    } finally {
      opened.session.stream.close();
    }
    if (!keep) { state.afterTime = 0; state.afterId = hex(new Uint8Array(proto.ID_LEN)); }
    save(state);
    console.log(`${total} message(s): ${stored} new, ${rejected} rejected`);
  },

  list() {
    const state = load();
    const wantRequests = rest.includes("--requests");
    for (const m of state.inbox || []) {
      if (Boolean(m.isRequest) === wantRequests) describe(state, m, "inbox");
    }
    if (rest.includes("--sent")) for (const m of state.sent || []) describe(state, m, "sent");
  },

  read() {
    const state = load();
    const wanted = rest[0].toLowerCase();
    const match = [...(state.inbox || []).map(m => [m, "inbox"]), ...(state.sent || []).map(m => [m, "sent"])]
      .find(([m]) => m.id.startsWith(wanted));
    if (!match) throw new proto.SmolError(`no message matching ${wanted}`);
    describe(state, match[0], match[1], true);
  },

  async rotate() {
    const state = load();
    const old = currentIdentity(state);
    const master = unhex(state.master);
    const account = proto.parseAddress(`${state.account.user}@${state.account.host}` +
      (state.account.port === proto.DEFAULT_PORT ? "" : `:${state.account.port}`));
    const nextIndex = (state.rotations ?? 0) + 1;
    const freshSeed = proto.identitySeed(master, nextIndex);
    const fresh = proto.identityFromSeed(freshSeed);
    const cert = proto.makeCert(account.user, old, freshSeed);
    const opened = await open(account.host, account.port, { requirePin: true });
    try {
      await proto.registerOp(opened.session, opened.serverStatic, { username: account.user, identity: fresh, cert });
    } finally {
      opened.session.stream.close();
    }
    state.rotations = nextIndex;
    save(state);
    console.log(`rotated ${account.short}`);
    console.log(`new key:     ${b32encode(fresh.publicKey)}`);
    console.log(`fingerprint: ${fingerprint(fresh.publicKey)}`);
  },
};

async function pushTokens(state) {
  if (!state.account) {
    console.error("warning: not registered; the set will be pushed with your first fetch");
    return 0;
  }
  const me = currentIdentity(state);
  const account = proto.parseAddress(`${state.account.user}@${state.account.host}` +
    (state.account.port === proto.DEFAULT_PORT ? "" : `:${state.account.port}`));
  const opened = await open(account.host, account.port, { requirePin: true });
  try {
    const { sync, tokens } = tokenSet(state);
    return await proto.authenticate(opened.session, opened.handshakeHash, account.user, me, { sync, tokens });
  } finally {
    opened.session.stream.close();
  }
}

function describe(state, row, box, verbose = false) {
  const id = row.id;
  try {
    const opened = proto.unseal(identities(state), unhex(row.envelope));
    const { fields, body } = proto.parseFrontmatter(new TextDecoder().decode(opened.body));
    const who = box === "sent" ? row.recipient
      : Object.entries(state.contacts || {}).find(([, c]) => c.key === b32encode(opened.sender))?.[0]
        ?? `<${b32encode(opened.sender).slice(0, 20)}…>`;
    const stamp = new Date((opened.time) * 1000).toISOString().slice(0, 16).replace("T", " ");
    if (!verbose) return console.log(
      `${id.slice(0, 8)}  ${stamp}  ${who.padEnd(28).slice(0, 28)}  ${fields.subject || "(no subject)"}`);
    console.log(`id:      ${id}`);
    console.log(`${box === "sent" ? "to" : "from"}:      ${who}`);
    console.log(`key:     ${b32encode(opened.sender)}`);
    console.log(`date:    ${new Date(opened.time * 1000).toISOString()}`);
    for (const [k, v] of Object.entries(fields)) if (k !== "accept") console.log(`${k}: ${v}`);
    console.log("signature verified\n");
    console.log(body);
  } catch (err) {
    if (verbose) throw err;
    console.log(`${id.slice(0, 8)}  <unreadable: ${err.message}>`);
  }
}

function cryptoRandom(n) {
  const out = new Uint8Array(n);
  crypto.getRandomValues(out);
  return out;
}

if (!commands[command]) {
  console.error(`unknown command: ${command}`);
  process.exit(2);
}
await commands[command]();