import org.bouncycastle.jcajce.provider.digest.SHA3;
import org.bouncycastle.pqc.crypto.mldsa.MLDSAParameters;
import org.bouncycastle.pqc.crypto.mldsa.MLDSAPrivateKeyParameters;
import org.bouncycastle.pqc.crypto.mldsa.MLDSAPublicKeyParameters;
import org.bouncycastle.pqc.crypto.mldsa.MLDSASigner;

import java.nio.charset.StandardCharsets;
import java.util.Arrays;

/** Probe ML-DSA-87 seed keygen + signing round-trip against noble vectors. */
public class MlDsaParityProbe {
    static String hex(byte[] b) {
        StringBuilder sb = new StringBuilder();
        for (byte x : b) sb.append(String.format("%02x", x));
        return sb.toString();
    }

    public static void main(String[] args) throws Exception {
        byte[] full = new byte[64];
        for (int i = 0; i < 64; i++) full[i] = (byte) i;
        SHA3.DigestSHA3 d = new SHA3.DigestSHA3(256);
        byte[] xi = d.digest(full);
        System.out.println("xi=" + hex(xi));

        MLDSAPrivateKeyParameters priv = new MLDSAPrivateKeyParameters(MLDSAParameters.ml_dsa_87, xi);
        byte[] pub = priv.getPublicKey();
        System.out.println("pubLen=" + pub.length);
        System.out.println("pubHead=" + hex(Arrays.copyOfRange(pub, 0, 32)));
        System.out.println("pubTail=" + hex(Arrays.copyOfRange(pub, pub.length - 32, pub.length)));
        System.out.println("encLen=" + priv.getEncoded().length);

        byte[] msg = "sarmad-parity".getBytes(StandardCharsets.UTF_8);
        MLDSASigner signer = new MLDSASigner();
        signer.init(true, priv);
        signer.update(msg, 0, msg.length);
        byte[] sig = signer.generateSignature();

        MLDSASigner verifier = new MLDSASigner();
        verifier.init(false, new MLDSAPublicKeyParameters(MLDSAParameters.ml_dsa_87, pub));
        verifier.update(msg, 0, msg.length);
        System.out.println("freshVerify=" + verifier.verifySignature(sig));

        MLDSAPrivateKeyParameters reloaded = new MLDSAPrivateKeyParameters(
                MLDSAParameters.ml_dsa_87, priv.getEncoded(), priv.getPublicKeyParameters());
        signer = new MLDSASigner();
        signer.init(true, reloaded);
        signer.update(msg, 0, msg.length);
        byte[] sig2 = signer.generateSignature();
        verifier = new MLDSASigner();
        verifier.init(false, new MLDSAPublicKeyParameters(MLDSAParameters.ml_dsa_87, pub));
        verifier.update(msg, 0, msg.length);
        System.out.println("reloadVerify=" + verifier.verifySignature(sig2));

        MLDSAPrivateKeyParameters fromSeed = new MLDSAPrivateKeyParameters(MLDSAParameters.ml_dsa_87, xi);
        signer = new MLDSASigner();
        signer.init(true, fromSeed);
        signer.update(msg, 0, msg.length);
        byte[] sig3 = signer.generateSignature();
        verifier = new MLDSASigner();
        verifier.init(false, new MLDSAPublicKeyParameters(MLDSAParameters.ml_dsa_87, pub));
        verifier.update(msg, 0, msg.length);
        System.out.println("seedReloadVerify=" + verifier.verifySignature(sig3));
        System.out.println("pubMatchSeed=" + Arrays.equals(pub, fromSeed.getPublicKey()));
    }
}
