#!/usr/bin/env bun
/**
 * Shade 4.0 GA — combined soak harness.
 *
 * Runs the ratchet stack under sustained ping-pong load for a configurable
 * duration. Designed to be wrapped in `systemd-run` / `nohup` / a Gitea
 * scheduled job and left running for ≥ 2 weeks before tagging GA-stable.
 *
 *   Usage:
 *     bun run scripts/soak.ts                 # default 1h
 *     bun run scripts/soak.ts --hours 168     # 1 week
 *     bun run scripts/soak.ts --hours 336     # 2 weeks (V4.0 §Soak)
 *     bun run scripts/soak.ts --hours 0.05    # 3 minute smoke (CI gate)
 */
import { ShadeSessionManager } from '../packages/shade-core/src/index.js';
import { SubtleCryptoProvider, MemoryStorage } from '../packages/shade-crypto-web/src/index.js';

interface Args {
  hours: number;
  pairs: number;
  reportSeconds: number;
}

function parseArgs(argv: string[]): Args {
  const out: Args = { hours: 1, pairs: 4, reportSeconds: 60 };
  for (let i = 0; i < argv.length; i++) {
    const a = argv[i];
    if (a === '--hours') out.hours = parseFloat(argv[++i]!);
    else if (a === '--pairs') out.pairs = parseInt(argv[++i]!, 10);
    else if (a === '--report-seconds') out.reportSeconds = parseInt(argv[++i]!, 10);
    else if (a === '--help') {
      console.log('Usage: bun run scripts/soak.ts [--hours N] [--pairs N] [--report-seconds N]');
      process.exit(0);
    }
  }
  return out;
}

interface Counters {
  ratchetSent: number;
  ratchetRecv: number;
  ratchetErrs: number;
  errors: { workload: string; message: string; ts: number }[];
}

function newCounters(): Counters {
  return {
    ratchetSent: 0,
    ratchetRecv: 0,
    ratchetErrs: 0,
    errors: [],
  };
}

interface Pair {
  id: number;
  aliceMgr: ShadeSessionManager;
  bobMgr: ShadeSessionManager;
  aliceAddr: string;
  bobAddr: string;
}

async function setupPair(id: number): Promise<Pair> {
  const crypto = new SubtleCryptoProvider();
  const aliceMgr = new ShadeSessionManager(crypto, new MemoryStorage());
  const bobMgr = new ShadeSessionManager(crypto, new MemoryStorage());
  await aliceMgr.initialize();
  await bobMgr.initialize();

  const otpks = await bobMgr.generateOneTimePreKeys(5);
  const bundle = await bobMgr.createPreKeyBundle();
  bundle.oneTimePreKey = { keyId: otpks[0]!.keyId, publicKey: otpks[0]!.keyPair.publicKey };

  const aliceAddr = `alice-${id}`;
  const bobAddr = `bob-${id}`;
  await aliceMgr.initSessionFromBundle(bobAddr, bundle);

  // Establish the receiver session by sending a first message and decrypting.
  const init = await aliceMgr.encrypt(bobAddr, 'init');
  await bobMgr.decrypt(aliceAddr, init);

  return { id, aliceMgr, bobMgr, aliceAddr, bobAddr };
}

async function ratchetPingPong(pair: Pair, counters: Counters, abort: AbortSignal) {
  const { aliceMgr, bobMgr, aliceAddr, bobAddr } = pair;
  let i = 0;
  while (!abort.aborted) {
    try {
      const env = await aliceMgr.encrypt(bobAddr, `msg-${i++}`);
      counters.ratchetSent++;
      await bobMgr.decrypt(aliceAddr, env);
      counters.ratchetRecv++;
      const reply = await bobMgr.encrypt(aliceAddr, `re-${i}`);
      counters.ratchetSent++;
      await aliceMgr.decrypt(bobAddr, reply);
      counters.ratchetRecv++;
      if ((i & 0xff) === 0) await new Promise((r) => setTimeout(r, 0));
    } catch (err) {
      counters.ratchetErrs++;
      counters.errors.push({
        workload: `pair-${pair.id}`,
        message: (err as Error).message,
        ts: Date.now(),
      });
      if (counters.ratchetErrs > 8) break;
    }
  }
}

async function main() {
  const args = parseArgs(process.argv.slice(2));
  const totalMs = Math.max(60_000, Math.round(args.hours * 3_600_000));
  const ac = new AbortController();
  const counters = newCounters();
  const start = Date.now();

  console.log(`[soak] start — ${args.pairs} pairs, ${args.hours}h, report every ${args.reportSeconds}s`);

  const pairs = await Promise.all(Array.from({ length: args.pairs }, (_, i) => setupPair(i)));

  setTimeout(() => ac.abort(), totalMs).unref();
  process.on('SIGINT', () => ac.abort());
  process.on('SIGTERM', () => ac.abort());

  const reporter = setInterval(() => {
    const elapsedS = ((Date.now() - start) / 1000).toFixed(0);
    const rps = (counters.ratchetSent / Math.max(1, (Date.now() - start) / 1000)).toFixed(1);
    console.log(
      `[soak] t=${elapsedS}s sent=${counters.ratchetSent} recv=${counters.ratchetRecv} errs=${counters.ratchetErrs} (${rps} sent/s)`,
    );
    if (counters.errors.length > 32) counters.errors.splice(0, counters.errors.length - 32);
  }, args.reportSeconds * 1000);
  reporter.unref();

  await Promise.all(pairs.map((p) => ratchetPingPong(p, counters, ac.signal)));

  clearInterval(reporter);

  const elapsedS = ((Date.now() - start) / 1000).toFixed(0);
  console.log(`\n[soak] done in ${elapsedS}s`);
  console.log(`  ratchet sent  ${counters.ratchetSent}`);
  console.log(`  ratchet recv  ${counters.ratchetRecv}`);
  console.log(`  ratchet errs  ${counters.ratchetErrs}`);

  if (counters.ratchetErrs > 0) {
    console.error('\nFAIL: ratchet errors recorded:');
    for (const e of counters.errors.slice(-10)) {
      console.error(`  [${new Date(e.ts).toISOString()}] ${e.workload}: ${e.message}`);
    }
    process.exit(1);
  }
  process.exit(0);
}

void main();