package meerkat.mixer;

import com.google.protobuf.InvalidProtocolBufferException;
import meerkat.mixer.network.BenesNetwork;
import meerkat.mixer.network.PermutationNetwork;
import meerkat.mixer.proofs.Mix2nizk;
import meerkat.protobuf.Crypto;
import meerkat.protobuf.Mixing;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Arrays;

/**
 * Created by Tzlil on 12/30/2015.
 * provide one operation - verify meerkat.mixProverVerifier.network output
 */
public final class MixVerifier {

    final static Logger logger = LoggerFactory.getLogger(MixVerifier.class);
    /**
     * constructor
     * @param verifier
     * @param mixerOutput
     * @param strict if true, check that the network structure matches our implementation of the Benes network exactly.
     *               if false, only checks that the result is a permutation.
     * @return true iff the meerkat.mixProverVerifier.network output is valid
     * @throws InvalidProtocolBufferException
     */
    public static boolean verifyTable(Mix2nizk.Verifier verifier, MixerOutput mixerOutput, boolean strict)
            throws InvalidProtocolBufferException {

        PermutationNetwork net = null;

        int layers = mixerOutput.getNumLayers();
        int n = 1 << mixerOutput.getLogN();

        if (strict) {
            net = new BenesNetwork(mixerOutput.getLogN());
        }

        //initialize outputUsed table
        // used to check BeneshNet validity
        boolean[][] outputUsed = new boolean[layers][n];
        for (boolean[] outputUsedLayer: outputUsed) {
            Arrays.fill(outputUsedLayer,false);
        }

        boolean[][] switchUsed = new boolean[layers][n / 2];
        for (boolean[] switchUsedLayer: switchUsed) {
            Arrays.fill(switchUsedLayer,false);
        }

        Mixing.Mix2Proof[][] Mix2Proofs = mixerOutput.getProofs();
        Crypto.RerandomizableEncryptedMessage[][] rerandomizableEncryptedMessages = mixerOutput.getEncryptedMessages();

        for (int i = 0; i < Mix2Proofs.length ; i++){
            for (int j = 0; j < Mix2Proofs[i].length ; j ++){
                int out0,out1,layer, switchIdx;

                Mixing.Mix2Proof zkp = Mix2Proofs[i][j];
                Mixing.Mix2Proof.Location location = zkp.getLocation();
                out0 = location.getOut0();
                out1 = location.getOut1();
                layer = location.getLayer();
                switchIdx = location.getSwitchIdx();

                if (strict) {
                    // Check that location is correct
                    int expectedOut0 = net.getInputIdxInNextLayer(layer, 2 * switchIdx);
                    if (out0 != expectedOut0) {
                        logger.error("Input {} in layer {} goes to {} instead of {}", 2 * switchIdx, layer, out0, expectedOut0);
                        return false;
                    }
                    int expectedOut1 = net.getInputIdxInNextLayer(layer, 2 * switchIdx + 1);
                    if (out0 != expectedOut0) {
                        logger.error("Input {} in layer {} goes to {} instead of {}", 2 * switchIdx + 1, layer, out1, expectedOut1);
                        return false;
                    }
                }

                // mark location in table
                if (switchUsed[layer][switchIdx]) {
                    logger.error("Switch {} in layer {} used twice!", switchIdx, layer);
                    return false;
                }
                switchUsed[layer][switchIdx] = true;

                if (outputUsed[layer][out0]) {
                    logger.error("Output {} in layer {} used twice!", out0, layer);
                    return false;
                }
                outputUsed[layer][out0] = true;

                if (outputUsed[layer][out1]) {
                    logger.error("Output {} in layer {} used twice!", out1, layer);
                    return false;
                }
                outputUsed[layer][out1] = true;

                // verify proof
                if(!verifier.verify(rerandomizableEncryptedMessages[layer][2 * switchIdx],
                        rerandomizableEncryptedMessages[layer][2 * switchIdx + 1],
                        rerandomizableEncryptedMessages[layer + 1][out0],
                        rerandomizableEncryptedMessages[layer + 1][out1],
                        zkp)) {
                    return false;
                }
            }
        }

        // verify all meerkat.mixProverVerifier.necessary locations for BenesNet were proved
        for (int layer = 0; layer < layers; ++layer) {
            for (int switchIdx = 0; switchIdx < n / 2; ++switchIdx) {
                if (!switchUsed[layer][switchIdx]) {
                    logger.error("Switch {} in layer {} was not used!", switchIdx, layer);
                    return false;
                }
            }
            for (int i = 0; i < n / 2; ++i) {
                if (!outputUsed[layer][i]) {
                    logger.error("Output {} in layer {} was not used!", i, layer);
                    return false;
                }
            }
        }
        return true;
    }
}
