package org.hyperledger.besu.evm.precompile;

// Backup ML-KEM-768 KAT generator (used only if the NIST ACVP fetch is unavailable).
// Produces vectors whose expected (c, k) are cryptographically trustworthy because
// each shared secret k is CONFIRMED by an independent decapsulation:
//   keygen -> (ek, dk); pick random m; (k, c) = Encaps_internal(ek, m);
//   k2 = Decaps(dk, c); assert k == k2  (ML-KEM correctness / KAT property).
// Emits tcId|ek_hex|m_hex|c_hex|k_hex, same format as the NIST ACVP file, for
// MlKemAcvpKat.java to drive through the precompile.

import org.bouncycastle.crypto.SecretWithEncapsulation;
import org.bouncycastle.pqc.crypto.mlkem.*;
import java.io.*;
import java.nio.file.*;
import java.security.SecureRandom;
import java.util.Arrays;

public class GenMlKemSelfKat {
  static String hex(byte[] b) {
    StringBuilder s = new StringBuilder();
    for (byte x : b) s.append(String.format("%02x", x));
    return s.toString();
  }

  public static void main(String[] a) throws Exception {
    String out = a.length > 0 ? a[0] : "mlkem768_acvp.txt";
    int n = a.length > 1 ? Integer.parseInt(a[1]) : 20;
    SecureRandom rnd = new SecureRandom();

    MLKEMKeyPairGenerator kg = new MLKEMKeyPairGenerator();
    kg.init(new MLKEMKeyGenerationParameters(rnd, MLKEMParameters.ml_kem_768));

    StringBuilder sb = new StringBuilder();
    sb.append("# BC-generated ML-KEM-768 encapsulation vectors, decapsulation-confirmed (ek|m|c|k)\n");
    int ok = 0;
    for (int i = 0; i < n; i++) {
      org.bouncycastle.crypto.AsymmetricCipherKeyPair kp = kg.generateKeyPair();
      MLKEMPublicKeyParameters pub = (MLKEMPublicKeyParameters) kp.getPublic();
      MLKEMPrivateKeyParameters priv = (MLKEMPrivateKeyParameters) kp.getPrivate();

      byte[] m = new byte[32];
      rnd.nextBytes(m);

      MLKEMGenerator gen = new MLKEMGenerator(rnd);
      SecretWithEncapsulation enc = gen.internalGenerateEncapsulated(pub, m);
      byte[] k = enc.getSecret();
      byte[] c = enc.getEncapsulation();

      // independent confirmation of k via decapsulation
      MLKEMExtractor ext = new MLKEMExtractor(priv);
      byte[] k2 = ext.extractSecret(c);
      if (!Arrays.equals(k, k2)) {
        System.out.println("DECAPS_MISMATCH at " + i);
        System.exit(4);
      }
      ok++;
      sb.append(i).append("|").append(hex(pub.getEncoded())).append("|").append(hex(m))
          .append("|").append(hex(c)).append("|").append(hex(k)).append("\n");
    }
    Files.write(Paths.get(out), sb.toString().getBytes());
    System.out.println("GenMlKemSelfKat: wrote " + ok + " decaps-confirmed vectors -> " + out);
  }
}