import java.security.*;
import java.security.spec.*;
import java.util.function.*;
import javax.crypto.*;

/** Operations per second, single thread, simple time-boxed loops. Indicative only: this is not JMH. */
public class PqcSpeed {
    interface Op { void run() throws Exception; }
    static volatile Object sink;

    static void bench(String name, Op op) throws Exception {
        long end = System.nanoTime() + 1_500_000_000L;
        while (System.nanoTime() < end) op.run();            // warmup
        long n = 0, t0 = System.nanoTime(); end = t0 + 2_000_000_000L;
        while (System.nanoTime() < end) { op.run(); n++; }
        double secs = (System.nanoTime() - t0) / 1e9;
        System.out.printf("%-34s %10.0f ops/s%n", name, n / secs);
    }

    public static void main(String[] a) throws Exception {
        System.out.println("java.version=" + System.getProperty("java.version"));
        byte[] msg = new byte[256];

        KeyPairGenerator kg = KeyPairGenerator.getInstance("ML-KEM-768");
        KeyPair kk = kg.generateKeyPair();
        KEM kem = KEM.getInstance("ML-KEM");
        KEM.Encapsulator enc = kem.newEncapsulator(kk.getPublic());
        byte[] ct = enc.encapsulate().encapsulation();
        KEM.Decapsulator dec = kem.newDecapsulator(kk.getPrivate());
        bench("ML-KEM-768 keygen", () -> sink = kg.generateKeyPair());
        bench("ML-KEM-768 encapsulate", () -> sink = kem.newEncapsulator(kk.getPublic()).encapsulate());
        bench("ML-KEM-768 decapsulate", () -> sink = dec.decapsulate(ct));

        KeyPairGenerator dg = KeyPairGenerator.getInstance("ML-DSA-65");
        KeyPair dk = dg.generateKeyPair();
        Signature ds = Signature.getInstance("ML-DSA-65");
        ds.initSign(dk.getPrivate()); ds.update(msg); byte[] dsig = ds.sign();
        bench("ML-DSA-65 keygen", () -> sink = dg.generateKeyPair());
        bench("ML-DSA-65 sign", () -> { Signature s = Signature.getInstance("ML-DSA-65"); s.initSign(dk.getPrivate()); s.update(msg); sink = s.sign(); });
        bench("ML-DSA-65 verify", () -> { Signature s = Signature.getInstance("ML-DSA-65"); s.initVerify(dk.getPublic()); s.update(msg); sink = s.verify(dsig); });

        KeyPairGenerator rg = KeyPairGenerator.getInstance("RSA"); rg.initialize(2048);
        KeyPair rk = rg.generateKeyPair();
        Signature rs = Signature.getInstance("SHA256withRSA"); rs.initSign(rk.getPrivate()); rs.update(msg); byte[] rsig = rs.sign();
        bench("RSA-2048 sign", () -> { Signature s = Signature.getInstance("SHA256withRSA"); s.initSign(rk.getPrivate()); s.update(msg); sink = s.sign(); });
        bench("RSA-2048 verify", () -> { Signature s = Signature.getInstance("SHA256withRSA"); s.initVerify(rk.getPublic()); s.update(msg); sink = s.verify(rsig); });

        KeyPairGenerator eg = KeyPairGenerator.getInstance("EC"); eg.initialize(new ECGenParameterSpec("secp256r1"));
        KeyPair ek = eg.generateKeyPair();
        Signature es = Signature.getInstance("SHA256withECDSA"); es.initSign(ek.getPrivate()); es.update(msg); byte[] esig = es.sign();
        bench("ECDSA P-256 sign", () -> { Signature s = Signature.getInstance("SHA256withECDSA"); s.initSign(ek.getPrivate()); s.update(msg); sink = s.sign(); });
        bench("ECDSA P-256 verify", () -> { Signature s = Signature.getInstance("SHA256withECDSA"); s.initVerify(ek.getPublic()); s.update(msg); sink = s.verify(esig); });

        KeyPairGenerator xg = KeyPairGenerator.getInstance("X25519");
        KeyPair xa = xg.generateKeyPair(), xb = xg.generateKeyPair();
        bench("X25519 keygen + agreement", () -> { KeyAgreement k = KeyAgreement.getInstance("X25519"); k.init(xa.getPrivate()); k.doPhase(xb.getPublic(), true); sink = k.generateSecret(); });
    }
}