package meerkat.mixer.proofs.generic;

import com.google.protobuf.Message;
import meerkat.mixer.proofs.Concatenator;
import meerkat.mixer.proofs.SigmaProtocol;

import java.math.BigInteger;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;

/**
 * Transform a Sigma protocol into a NIZK using Fiat-Shamir
 * We use SHA-256 explicitly to make it easier to verify in other codebases
 */
public class SigmaFiatShamir<NIZKMsgType, FirstMsgType extends Message, FinalMessageType> {
    final static String DIGEST_ALG = "SHA-256";
    final MessageDigest md;
    final Concatenator.Pair<NIZKMsgType, FirstMsgType, FinalMessageType> concat;

    public SigmaFiatShamir(Concatenator.Pair<NIZKMsgType, FirstMsgType, FinalMessageType> concat) {
        this.concat = concat;
        try {
            md = MessageDigest.getInstance(DIGEST_ALG);
        } catch (NoSuchAlgorithmException e) {
            // Should never happen
            throw new RuntimeException("Error in instantiating " + DIGEST_ALG + " digest", e);
        }
    }

    /**
     * Fiat-Shamir heuristic
     * @param input - protobuf contains all parameters from the first step of the current proof
     * @return randomOracle.hash(input)
     */
    public BigInteger hash(Message input) {
        md.reset();

        byte[] arr = input.toByteArray();
        byte[] digest = md.digest(arr);
        return new BigInteger(1,digest);
    }

    public NIZKMsgType generateNizk(SigmaProtocol.Prover<FirstMsgType, FinalMessageType> prover) {
        FirstMsgType firstMessage = prover.getFirstMessage();
        BigInteger challenge = hash(firstMessage);
        FinalMessageType finalMessage = prover.getFinalMessage(challenge);
        return concat.concatenate(firstMessage, finalMessage);
    }

    public boolean verifyNizk(NIZKMsgType NIZK, SigmaProtocol.Verifier<FirstMsgType, FinalMessageType> verifier) {
        FirstMsgType firstMessage = concat.getMsg1(NIZK);
        BigInteger challenge = hash(firstMessage);
        return verifier.verify(firstMessage, challenge, concat.getMsg2(NIZK));
    }
}
