package org.hyperledger.besu.evm.precompile;

// AERE ML-KEM-768 precompile (0x0AE6) KAT harness.
//
// Drives the REAL MLKEM768PrecompiledContract.computePrecompile over NIST ACVP
// ML-KEM-encapDecap-FIPS203 encapsulation vectors. Each vectors-file line is:
//   tcId|ek_hex|m_hex|expected_c_hex|expected_k_hex
// The harness assembles the precompile input (ek || m), runs the precompile, and
// asserts the 1120-byte output equals (expected_c || expected_k) byte-for-byte.
// It also re-runs each vector to confirm the encapsulation is deterministic.
//
// Writes a JSON results file (path in argv[1]) mirroring nethermind-kat-results.

import org.apache.tuweni.bytes.Bytes;
import org.hyperledger.besu.evm.gascalculator.CancunGasCalculator;
import java.io.*;
import java.nio.file.*;
import java.util.*;

public class MlKemAcvpKat {

  static byte[] hex(String s) {
    s = s.trim();
    if (s.startsWith("0x")) s = s.substring(2);
    if (s.isEmpty()) return new byte[0];
    byte[] b = new byte[s.length() / 2];
    for (int i = 0; i < b.length; i++)
      b[i] = (byte) Integer.parseInt(s.substring(2 * i, 2 * i + 2), 16);
    return b;
  }

  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 vectorsPath = a.length > 0 ? a[0] : "mlkem768_acvp.txt";
    String outPath = a.length > 1 ? a[1] : "kat-results-mlkem.json";

    MLKEM768PrecompiledContract pc = new MLKEM768PrecompiledContract(new CancunGasCalculator());

    int passed = 0, failed = 0, total = 0;
    StringBuilder results = new StringBuilder();
    results.append("{\n  \"precompile\": \"ML-KEM-768 (FIPS 203) deterministic encapsulate\",\n");
    results.append("  \"address\": \"0x0000000000000000000000000000000000000ae6\",\n");
    results.append("  \"source\": \"NIST ACVP ML-KEM-encapDecap-FIPS203 (encapsulation AFT)\",\n");
    results.append("  \"results\": [\n");

    List<String> lines = Files.readAllLines(Paths.get(vectorsPath));
    boolean first = true;
    for (String raw : lines) {
      String line = raw.trim();
      if (line.isEmpty() || line.startsWith("#")) continue;
      String[] p = line.split("\\|");
      if (p.length != 5) continue;
      total++;
      String tcId = p[0];
      byte[] ek = hex(p[1]);
      byte[] m = hex(p[2]);
      byte[] wantC = hex(p[3]);
      byte[] wantK = hex(p[4]);

      byte[] input = new byte[ek.length + m.length];
      System.arraycopy(ek, 0, input, 0, ek.length);
      System.arraycopy(m, 0, input, ek.length, m.length);

      Bytes out1 = pc.computePrecompile(Bytes.wrap(input), null).output();
      Bytes out2 = pc.computePrecompile(Bytes.wrap(input), null).output(); // determinism check

      byte[] want = new byte[wantC.length + wantK.length];
      System.arraycopy(wantC, 0, want, 0, wantC.length);
      System.arraycopy(wantK, 0, want, wantC.length, wantK.length);

      boolean match = out1 != null && Arrays.equals(out1.toArrayUnsafe(), want);
      boolean deterministic = out1 != null && out2 != null
          && Arrays.equals(out1.toArrayUnsafe(), out2.toArrayUnsafe());
      boolean ok = match && deterministic;
      if (ok) passed++; else failed++;

      if (!first) results.append(",\n");
      first = false;
      results.append("    {\"tcId\": \"").append(tcId).append("\", \"ekLen\": ").append(ek.length)
          .append(", \"ctLen\": ").append(out1 == null ? 0 : out1.size())
          .append(", \"expectedMatch\": ").append(match)
          .append(", \"deterministic\": ").append(deterministic)
          .append(", \"status\": \"").append(ok ? "PASS" : "FAIL").append("\"}");
      System.out.println("MLKEM768 tc" + tcId + " match=" + match + " deterministic=" + deterministic
          + " -> " + (ok ? "PASS" : "FAIL"));
    }

    results.append("\n  ],\n  \"passed\": ").append(passed).append(",\n  \"failed\": ")
        .append(failed).append(",\n  \"total\": ").append(total).append("\n}\n");
    Files.write(Paths.get(outPath), results.toString().getBytes());

    System.out.println("MLKEM768 KAT: passed=" + passed + " failed=" + failed + " total=" + total);
    if (failed != 0 || total == 0) {
      System.out.println("MLKEM768_KAT_FAILED");
      System.exit(3);
    }
    System.out.println("MLKEM768_KAT_OK");
  }
}