import org.bouncycastle.jcajce.provider.digest.SHA3;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyGenerationParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumKeyPairGenerator;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumParameters;
import org.bouncycastle.pqc.crypto.crystals.dilithium.DilithiumPublicKeyParameters;

import java.security.SecureRandom;
import java.util.Arrays;

/** Test whether Dilithium keygen with fixed-first-32 RNG matches noble ML-DSA. */
public class DilithiumSeedProbe {
    public static void main(String[] args) throws Exception {
        byte[] fullSeed = new byte[64];
        for (int i = 0; i < 64; i++) fullSeed[i] = (byte) i;
        byte[] sha3Seed = new SHA3.Digest256().digest(fullSeed);

        SecureRandom fixed = new SecureRandom() {
            int calls = 0;
            @Override public void nextBytes(byte[] bytes) {
                calls++;
                Arrays.fill(bytes, (byte) 0);
                if (calls == 1) {
                    System.arraycopy(sha3Seed, 0, bytes, 0, Math.min(32, bytes.length));
                    System.out.println("call1 len=" + bytes.length + " seed=" + hex(Arrays.copyOf(bytes, Math.min(32, bytes.length))));
                } else {
                    System.out.println("call" + calls + " len=" + bytes.length);
                }
            }
            @Override public byte[] generateSeed(int numBytes) {
                byte[] b = new byte[numBytes];
                nextBytes(b);
                return b;
            }
        };

        DilithiumKeyPairGenerator keyGen = new DilithiumKeyPairGenerator();
        keyGen.init(new DilithiumKeyGenerationParameters(fixed, DilithiumParameters.dilithium5));
        var pair = keyGen.generateKeyPair();
        byte[] pub = ((DilithiumPublicKeyParameters) pair.getPublic()).getEncoded();
        System.out.println("pubHead=" + hex(Arrays.copyOf(pub, 32)));
        System.out.println("pubLen=" + pub.length);
    }

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