import { performance } from "node:perf_hooks";

const target = process.env.ZDS_BENCH_TARGET;
const accountCount = Number(process.env.ZDS_BENCH_ACCOUNTS ?? "100");
const matrix = parseMatrix(
  process.env.ZDS_BENCH_MATRIX ??
    "65536:20,100;3145728:1,3,5,10,20,40,60,100;20971520:1,3,5,10",
);
const timeoutMs = Number(process.env.ZDS_BENCH_TIMEOUT_MS ?? "60000");

if (!target) throw new Error("ZDS_BENCH_TARGET is required");

const sessions = await loginAccounts();
const results = [];

for (const { bytes, concurrencies } of matrix) {
  const payload = Buffer.alloc(bytes, 0x5a);
  for (const concurrency of concurrencies) {
    const result = await runStage(payload, concurrency);
    results.push(result);
    console.log(JSON.stringify(result));
  }
}

console.log("\nsize\tusers\tok\tfail\tpublish p50\tp95\tp99\tmax\tprobe p95");
for (const row of results) {
  console.log(
    [
      formatBytes(row.bytes),
      row.concurrency,
      row.ok,
      row.failed,
      formatMs(row.publish.p50),
      formatMs(row.publish.p95),
      formatMs(row.publish.p99),
      formatMs(row.publish.max),
      formatMs(row.probe.p95),
    ].join("\t"),
  );
}

if (results.some((row) => row.failed > 0)) process.exitCode = 1;

async function loginAccounts() {
  const output = new Array(accountCount);
  await parallel(accountCount, 20, async (index) => {
    const number = index + 1;
    const handle = `bench${String(number).padStart(3, "0")}.test`;
    const response = await request("/xrpc/com.atproto.server.createSession", {
      method: "POST",
      headers: { "content-type": "application/json" },
      body: JSON.stringify({ identifier: handle, password: "password" }),
    });
    output[index] = {
      did: response.did,
      token: response.accessJwt,
    };
  });
  return output;
}

async function runStage(payload, concurrency) {
  const publish = [];
  const upload = [];
  const write = [];
  const verify = [];
  const probe = [];
  const errors = [];
  let probing = true;
  let probeIndex = 0;

  const probeTask = (async () => {
    while (probing) {
      const session = sessions[probeIndex++ % sessions.length];
      const started = performance.now();
      try {
        await request("/xrpc/com.atproto.server.getSession", {
          headers: { authorization: `Bearer ${session.token}` },
        });
        probe.push(performance.now() - started);
      } catch (error) {
        probe.push(performance.now() - started);
        errors.push(`probe: ${error.message}`);
      }
      await sleep(25);
    }
  })();

  const started = performance.now();
  await Promise.all(
    sessions.slice(0, concurrency).map(async (session, index) => {
      const operationStarted = performance.now();
      try {
        const uploadStarted = performance.now();
        const uploaded = await request("/xrpc/com.atproto.repo.uploadBlob", {
          method: "POST",
          headers: {
            authorization: `Bearer ${session.token}`,
            "content-type": "audio/mpeg",
          },
          body: payload,
        });
        upload.push(performance.now() - uploadStarted);

        const rkey = `bench-${Date.now()}-${index}`;
        const writeStarted = performance.now();
        await request("/xrpc/com.atproto.repo.createRecord", {
          method: "POST",
          headers: {
            authorization: `Bearer ${session.token}`,
            "content-type": "application/json",
          },
          body: JSON.stringify({
            repo: session.did,
            collection: "com.example.bench.record",
            rkey,
            validate: false,
            record: {
              $type: "com.example.bench.record",
              blob: uploaded.blob,
            },
          }),
        });
        write.push(performance.now() - writeStarted);

        const verifyStarted = performance.now();
        await request(
          `/xrpc/com.atproto.repo.getRecord?repo=${encodeURIComponent(session.did)}&collection=com.example.bench.record&rkey=${encodeURIComponent(rkey)}`,
          { headers: { authorization: `Bearer ${session.token}` } },
        );
        verify.push(performance.now() - verifyStarted);
        publish.push(performance.now() - operationStarted);
      } catch (error) {
        publish.push(performance.now() - operationStarted);
        errors.push(`user ${index + 1}: ${error.message}`);
      }
    }),
  );
  const elapsedMs = performance.now() - started;
  probing = false;
  await probeTask;

  return {
    bytes: payload.length,
    concurrency,
    elapsedMs,
    ok: concurrency - errors.filter((error) => error.startsWith("user ")).length,
    failed: errors.filter((error) => error.startsWith("user ")).length,
    publish: distribution(publish),
    upload: distribution(upload),
    write: distribution(write),
    verify: distribution(verify),
    probe: distribution(probe),
    errors: errors.slice(0, 10),
  };
}

async function request(path, init = {}) {
  const response = await fetch(`${target}${path}`, {
    ...init,
    signal: AbortSignal.timeout(timeoutMs),
  });
  const text = await response.text();
  if (!response.ok) {
    throw new Error(`${response.status} ${text.slice(0, 200)}`);
  }
  return text ? JSON.parse(text) : {};
}

async function parallel(count, concurrency, operation) {
  let next = 0;
  await Promise.all(
    Array.from({ length: Math.min(count, concurrency) }, async () => {
      while (true) {
        const index = next++;
        if (index >= count) return;
        await operation(index);
      }
    }),
  );
}

function distribution(values) {
  if (values.length === 0) {
    return { count: 0, p50: 0, p95: 0, p99: 0, max: 0 };
  }
  const sorted = [...values].sort((a, b) => a - b);
  return {
    count: sorted.length,
    p50: percentile(sorted, 0.5),
    p95: percentile(sorted, 0.95),
    p99: percentile(sorted, 0.99),
    max: sorted.at(-1),
  };
}

function percentile(sorted, quantile) {
  return sorted[Math.min(sorted.length - 1, Math.ceil(sorted.length * quantile) - 1)];
}

function parseMatrix(value) {
  return value.split(";").map((stage) => {
    const [bytes, concurrencies] = stage.split(":");
    return {
      bytes: Number(bytes),
      concurrencies: concurrencies.split(",").map(Number),
    };
  });
}

function formatBytes(value) {
  return value >= 1024 * 1024
    ? `${(value / 1024 / 1024).toFixed(0)}MiB`
    : `${(value / 1024).toFixed(0)}KiB`;
}

function formatMs(value) {
  return `${value.toFixed(1)}ms`;
}

function sleep(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}