#!/usr/bin/env node
/**
 * Pub→sub round trip through the cluster (two connections).
 * Usage: NATS_URL=... node latency.mjs [count] [payloadBytes] [publishers] [ping|flood|reconnect]
 * ping       = sequential publish-wait on persistent sockets (one-message RTT)
 * flood      = publish the batch then drain (queueing under burst)
 * reconnect  = connect, one publish, wait, close — measures handshake tax
 */
import { connect, headers } from "nats";

const url = process.env.NATS_URL || "nats://10.10.10.21:4222";
const count = Number(process.argv[2] || 5000);
const size = Number(process.argv[3] || 128);
const pubs = Number(process.argv[4] || 1);
const mode = process.argv[5] || "flood";
const servers = url.split(",").map((s) => s.trim());
const subject = `bench.lat.${process.pid}`;
const payload = new Uint8Array(size);

function pct(sorted, p) {
  if (!sorted.length) return 0;
  const i = Math.min(sorted.length - 1, Math.floor((p / 100) * sorted.length));
  return sorted[i];
}

const samples = [];
if (mode === "reconnect") {
  const subNc = await connect({ servers, name: "lat-sub" });
  let resolveOne = null;
  const sub = subNc.subscribe(subject, { max: count });
  const consume = (async () => {
    for await (const m of sub) {
      const sent = Number(m.headers?.get("t") || 0);
      samples.push(Number(process.hrtime.bigint() / 1000n) - sent);
      resolveOne?.();
    }
  })();
  await subNc.flush();
  for (let i = 0; i < count; i++) {
    const got = new Promise((r) => {
      resolveOne = r;
    });
    const pubNc = await connect({ servers, name: `lat-re-${i}` });
    const h = headers();
    h.set("t", String(process.hrtime.bigint() / 1000n));
    pubNc.publish(subject, payload, { headers: h });
    await pubNc.flush();
    await got;
    await pubNc.close();
  }
  await consume;
  await subNc.close();
} else if (mode === "ping") {
  const subNc = await connect({ servers, name: "lat-sub" });
  const pubNc = await connect({ servers, name: "lat-pub" });
  let resolveOne = null;
  const sub = subNc.subscribe(subject, { max: count });
  const consume = (async () => {
    for await (const m of sub) {
      const sent = Number(m.headers?.get("t") || 0);
      samples.push(Number(process.hrtime.bigint() / 1000n) - sent);
      resolveOne?.();
    }
  })();
  await subNc.flush();
  for (let i = 0; i < count; i++) {
    const got = new Promise((r) => {
      resolveOne = r;
    });
    const h = headers();
    h.set("t", String(process.hrtime.bigint() / 1000n));
    pubNc.publish(subject, payload, { headers: h });
    await got;
  }
  await consume;
  await pubNc.close();
  await subNc.close();
} else {
  const subNc = await connect({ servers, name: "lat-sub" });
  const sub = subNc.subscribe(subject, { max: count });
  const done = (async () => {
    for await (const m of sub) {
      const sent = Number(m.headers?.get("t") || 0);
      if (sent) samples.push(Number(process.hrtime.bigint() / 1000n) - sent);
    }
  })();
  await subNc.flush();
  const per = Math.ceil(count / pubs);
  const publishers = [];
  for (let p = 0; p < pubs; p++) {
    publishers.push(
      (async () => {
        const nc = await connect({ servers, name: `lat-pub-${p}` });
        const n = p === pubs - 1 ? count - per * (pubs - 1) : per;
        for (let i = 0; i < n; i++) {
          const h = headers();
          h.set("t", String(process.hrtime.bigint() / 1000n));
          nc.publish(subject, payload, { headers: h });
        }
        await nc.flush();
        await nc.close();
      })(),
    );
  }
  await Promise.all(publishers);
  await done;
  await subNc.close();
}

samples.sort((a, b) => a - b);
const sum = samples.reduce((a, b) => a + b, 0);
const us = (n) => `${(n / 1000).toFixed(3)}ms`;
console.log(
  JSON.stringify({
    count: samples.length,
    pubs,
    size,
    mode,
    min_us: samples[0],
    avg_us: Math.round(sum / samples.length),
    p50_us: pct(samples, 50),
    p90_us: pct(samples, 90),
    p99_us: pct(samples, 99),
    max_us: samples[samples.length - 1],
    min: us(samples[0]),
    avg: us(sum / samples.length),
    p50: us(pct(samples, 50)),
    p90: us(pct(samples, 90)),
    p99: us(pct(samples, 99)),
    max: us(samples[samples.length - 1]),
  }),
);