package org.hyperledger.besu.evm.precompile;

// AERE Falcon HashToPoint precompile (0x0AE7) KAT + vector generator.
//
// 1. Generates real Falcon-512 signatures with Bouncy Castle, extracts each
//    (nonce, message), and runs them through the REAL HashToPointPrecompiledContract.
// 2. Emits two files:
//      hashtopoint_vectors.txt : tcId|logn|nonce_hex|message_hex|precompile_coeffs_hex
//        (the coeffs are the precompile output; htp_reference.py recomputes them
//         independently with Python's stdlib SHAKE256 and must match)
//      falcon512_bench_vectors.txt : tcId|pk_hex|sm_hex
//        (real (pk, signed-message) pairs for the on-chain Falcon-verify benchmark
//         and the precompile-vs-in-EVM integration check)
// 3. Adds deterministic edge vectors (all-zero nonce, empty message, a logn=10 case).

import org.apache.tuweni.bytes.Bytes;
import org.bouncycastle.crypto.AsymmetricCipherKeyPair;
import org.bouncycastle.pqc.crypto.falcon.*;
import org.hyperledger.besu.evm.gascalculator.CancunGasCalculator;
import java.io.*;
import java.nio.file.*;
import java.security.SecureRandom;
import java.util.*;

public class HashToPointKat {

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

  static HashToPointPrecompiledContract PC =
      new HashToPointPrecompiledContract(new CancunGasCalculator());

  // Run the precompile for logn || nonce || message, return coeff bytes (2*n).
  static byte[] htp(int logn, byte[] nonce, byte[] message) {
    byte[] input = new byte[1 + nonce.length + message.length];
    input[0] = (byte) logn;
    System.arraycopy(nonce, 0, input, 1, nonce.length);
    System.arraycopy(message, 0, input, 1 + nonce.length, message.length);
    Bytes out = PC.computePrecompile(Bytes.wrap(input), null).output();
    return out == null ? new byte[0] : out.toArrayUnsafe();
  }

  public static void main(String[] a) throws Exception {
    String htpVec = a.length > 0 ? a[0] : "hashtopoint_vectors.txt";
    String benchVec = a.length > 1 ? a[1] : "falcon512_bench_vectors.txt";
    int nSigs = a.length > 2 ? Integer.parseInt(a[2]) : 8;

    SecureRandom rnd = new SecureRandom();
    StringBuilder htp = new StringBuilder();
    StringBuilder bench = new StringBuilder();
    int tc = 0;

    // ---- real Falcon-512 signatures ----
    FalconKeyPairGenerator kg = new FalconKeyPairGenerator();
    kg.init(new FalconKeyGenerationParameters(rnd, FalconParameters.falcon_512));
    for (int i = 0; i < nSigs; i++) {
      AsymmetricCipherKeyPair kp = kg.generateKeyPair();
      FalconPublicKeyParameters pub = (FalconPublicKeyParameters) kp.getPublic();
      FalconPrivateKeyParameters priv = (FalconPrivateKeyParameters) kp.getPrivate();
      byte[] message = ("AERE Falcon HashToPoint KAT vector #" + i + " (isolated fork)").getBytes();

      FalconSigner fs = new FalconSigner();
      fs.init(true, priv);
      byte[] bcSig = fs.generateSignature(message); // (0x30|logn) || nonce(40) || compressed
      byte[] nonce = Arrays.copyOfRange(bcSig, 1, 41);
      byte[] compressed = Arrays.copyOfRange(bcSig, 41, bcSig.length);

      // HashToPoint vector
      byte[] coeffs = htp(9, nonce, message);
      htp.append(tc).append("|9|").append(hex(nonce)).append("|").append(hex(message))
          .append("|").append(hex(coeffs)).append("\n");

      // Falcon bench vector: pk = 0x09 || h ; sm = sigLen(2) || nonce(40) || message || (0x29||compressed)
      byte[] h = pub.getH();
      byte[] pk = new byte[1 + h.length];
      pk[0] = 0x09;
      System.arraycopy(h, 0, pk, 1, h.length);
      byte[] sigBlock = new byte[1 + compressed.length];
      sigBlock[0] = (byte) 0x29; // 0x20 | logn(9)
      System.arraycopy(compressed, 0, sigBlock, 1, compressed.length);
      int sigLen = sigBlock.length;
      ByteArrayOutputStream sm = new ByteArrayOutputStream();
      sm.write((sigLen >> 8) & 0xff);
      sm.write(sigLen & 0xff);
      sm.write(nonce);
      sm.write(message);
      sm.write(sigBlock);
      bench.append(tc).append("|").append(hex(pk)).append("|").append(hex(sm.toByteArray())).append("\n");
      tc++;
    }

    // ---- deterministic edge vectors (HashToPoint only) ----
    byte[] zeroNonce = new byte[40];
    byte[][] edgeMsgs = new byte[][] {
      new byte[0],
      "abc".getBytes(),
      "The quick brown fox jumps over the lazy dog".getBytes()
    };
    for (byte[] msg : edgeMsgs) {
      byte[] c9 = htp(9, zeroNonce, msg);
      htp.append(tc).append("|9|").append(hex(zeroNonce)).append("|").append(hex(msg))
          .append("|").append(hex(c9)).append("\n");
      tc++;
    }
    // one Falcon-1024 sized (logn=10) HashToPoint vector
    byte[] c10 = htp(10, zeroNonce, "AERE Falcon-1024 HashToPoint".getBytes());
    htp.append(tc).append("|10|").append(hex(zeroNonce)).append("|")
        .append(hex("AERE Falcon-1024 HashToPoint".getBytes())).append("|").append(hex(c10)).append("\n");
    tc++;

    Files.write(Paths.get(htpVec), htp.toString().getBytes());
    Files.write(Paths.get(benchVec), bench.toString().getBytes());
    System.out.println("HashToPoint: wrote " + tc + " vectors to " + htpVec
        + " and " + nSigs + " Falcon bench vectors to " + benchVec);
  }
}