/**
 * Shade benchmarks.
 *
 * Run: bun bench/run.ts
 *
 * Output: console table + bench/results.md
 */
import { ShadeSessionManager, computeFingerprint } from '../packages/shade-core/src/index.js';
import { SubtleCryptoProvider, MemoryStorage } from '../packages/shade-crypto-web/src/index.js';
import {
  generateIdentityKeyPair,
  generateSignedPreKey,
  generateOneTimePreKeys,
  createPreKeyBundle,
  processPreKeyBundle,
  initSenderSession,
  ratchetEncrypt,
  ratchetDecrypt,
} from '../packages/shade-core/src/index.js';

const crypto = new SubtleCryptoProvider();

interface BenchResult {
  name: string;
  iterations: number;
  totalMs: number;
  perOpUs: number;
  opsPerSec: number;
}

const results: BenchResult[] = [];

async function bench(name: string, iterations: number, fn: () => Promise<void> | void) {
  // Warm up
  for (let i = 0; i < Math.min(10, iterations); i++) await fn();

  const start = performance.now();
  for (let i = 0; i < iterations; i++) await fn();
  const totalMs = performance.now() - start;

  const perOpUs = (totalMs * 1000) / iterations;
  const opsPerSec = (iterations / totalMs) * 1000;
  results.push({ name, iterations, totalMs, perOpUs, opsPerSec });
  console.log(`  ${name.padEnd(45)} ${perOpUs.toFixed(2).padStart(10)} µs/op  ${Math.round(opsPerSec).toLocaleString().padStart(10)} ops/sec`);
}

async function main() {
  console.log('=== Shade Benchmarks ===\n');

  // ─── Crypto primitives ─────────────────────────────────
  console.log('Crypto primitives:');

  await bench('X25519 keypair generation', 1000, async () => {
    await crypto.generateX25519KeyPair();
  });

  const a = await crypto.generateX25519KeyPair();
  const b = await crypto.generateX25519KeyPair();
  await bench('X25519 DH (shared secret)', 1000, async () => {
    await crypto.x25519(a.privateKey, b.publicKey);
  });

  await bench('Ed25519 keypair generation', 1000, async () => {
    await crypto.generateEd25519KeyPair();
  });

  const sigKp = await crypto.generateEd25519KeyPair();
  const msg = new TextEncoder().encode('test message');
  await bench('Ed25519 sign', 1000, async () => {
    await crypto.sign(sigKp.privateKey, msg);
  });

  const sig = await crypto.sign(sigKp.privateKey, msg);
  await bench('Ed25519 verify', 1000, async () => {
    await crypto.verify(sigKp.publicKey, msg, sig);
  });

  const aesKey = crypto.randomBytes(32);
  const plaintext = new TextEncoder().encode('a small message');
  await bench('AES-256-GCM encrypt (small)', 5000, async () => {
    await crypto.aesGcmEncrypt(aesKey, plaintext);
  });

  const enc = await crypto.aesGcmEncrypt(aesKey, plaintext);
  await bench('AES-256-GCM decrypt (small)', 5000, async () => {
    await crypto.aesGcmDecrypt(aesKey, enc.ciphertext, enc.nonce);
  });

  // ─── X3DH ──────────────────────────────────────────────
  console.log('\nX3DH handshake:');

  await bench('Generate identity keypair', 500, async () => {
    await generateIdentityKeyPair(crypto);
  });

  const bobIdentity = await generateIdentityKeyPair(crypto);
  await bench('Generate signed prekey', 500, async () => {
    await generateSignedPreKey(crypto, bobIdentity, 1);
  });

  const bobSpk = await generateSignedPreKey(crypto, bobIdentity, 1);
  const bundle = createPreKeyBundle(1, bobIdentity, bobSpk);

  await bench('Process prekey bundle (Alice X3DH)', 500, async () => {
    const aliceStorage = new MemoryStorage();
    const aliceIdentity = await generateIdentityKeyPair(crypto);
    await aliceStorage.saveIdentityKeyPair(aliceIdentity);
    await processPreKeyBundle(crypto, aliceStorage, bundle);
  });

  // ─── Double Ratchet ────────────────────────────────────
  console.log('\nDouble Ratchet:');

  // Set up a long-lived session for ratchet benchmarks
  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 bundle2 = await bobMgr.createPreKeyBundle();
  bundle2.oneTimePreKey = { keyId: otpks[0].keyId, publicKey: otpks[0].keyPair.publicKey };
  await aliceMgr.initSessionFromBundle('bob', bundle2);
  // Establish bidirectional session
  const env0 = await aliceMgr.encrypt('bob', 'init');
  await bobMgr.decrypt('alice', env0);
  const reply0 = await bobMgr.encrypt('alice', 'init reply');
  await aliceMgr.decrypt('bob', reply0);

  // Encrypt-only: needs careful counter management
  await bench('Encrypt message (no decrypt)', 500, async () => {
    await aliceMgr.encrypt('bob', 'hello world');
  });

  // Set up a fresh session for the roundtrip bench (Alice's chain is now far ahead)
  const alice2 = new ShadeSessionManager(crypto, new MemoryStorage());
  const bob2 = new ShadeSessionManager(crypto, new MemoryStorage());
  await alice2.initialize();
  await bob2.initialize();
  const otpks2 = await bob2.generateOneTimePreKeys(5);
  const bundle3 = await bob2.createPreKeyBundle();
  bundle3.oneTimePreKey = { keyId: otpks2[0].keyId, publicKey: otpks2[0].keyPair.publicKey };
  await alice2.initSessionFromBundle('bob', bundle3);
  const initEnv = await alice2.encrypt('bob', 'init');
  await bob2.decrypt('alice', initEnv);
  const initReply = await bob2.encrypt('alice', 'init reply');
  await alice2.decrypt('bob', initReply);

  await bench('Encrypt + decrypt roundtrip (in-sync)', 500, async () => {
    const env = await alice2.encrypt('bob', 'roundtrip');
    await bob2.decrypt('alice', env);
  });

  // ─── Fingerprint ───────────────────────────────────────
  console.log('\nFingerprint:');

  const sigKey = crypto.randomBytes(32);
  const dhKey = crypto.randomBytes(32);
  await bench('Compute fingerprint', 1000, async () => {
    await computeFingerprint(crypto, sigKey, dhKey);
  });

  // ─── Output to markdown ────────────────────────────────
  let md = '# Shade Benchmarks\n\n';
  md += `Generated: ${new Date().toISOString()}\n\n`;
  md += '| Operation | Iterations | µs/op | ops/sec |\n';
  md += '|-----------|------------|-------|--------|\n';
  for (const r of results) {
    md += `| ${r.name} | ${r.iterations} | ${r.perOpUs.toFixed(2)} | ${Math.round(r.opsPerSec).toLocaleString()} |\n`;
  }

  await Bun.write('bench/results.md', md);
  console.log('\nResults written to bench/results.md');
}

main().catch(console.error);