package meerkat.mixer;

import meerkat.protobuf.Crypto;
import meerkat.protobuf.Mixing;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;


/**
 * Created by Tzlil on 1/18/2016.
 * container for meerkat.mixProverVerifier.network.mix result.
 */
public class MixerOutput {
    private final Mixing.Mix2Proof[][] proofs;
    private final Crypto.RerandomizableEncryptedMessage[][] encryptedMessages;
    private final int logN;

    /**
     * constructor
     * @param logN log (base 2) of the number of votes
     * @param encryptedMessages at level 0 , contains the original encrypted votes
     *                          at each other level contains the re encrypted votes
     * @param proofs in each cell (level,switch) contains the match zero knowledge proof.
     *               We allow null proofs as a special case -- this is used to store the input to the
     *               first mixProverVerifier (as a 0-layer mix).
     */
    public MixerOutput(int logN, Mixing.Mix2Proof[][] proofs,
                       Crypto.RerandomizableEncryptedMessage[][] encryptedMessages) {
        this.proofs = proofs;
        this.encryptedMessages = encryptedMessages;
        this.logN = logN;
    }

    public Mixing.Mix2Proof[][] getProofs() {
        return proofs;
    }

    public Crypto.RerandomizableEncryptedMessage[][] getEncryptedMessages() {
        return encryptedMessages;
    }

    public int getLogN() {
        return logN;
    }

    public int getNumLayers() {
        return proofs == null ? 0 : proofs.length;
    }


    /**
     * Read from an InputStream. The data should have been written by {@link #writeDelimitedTo(OutputStream)}
     * @param in
     * @return
     * @throws IOException
     */
    public static MixerOutput parseDelimitedFrom(InputStream in) throws IOException {
        Mixing.MixBatchHeader header = Mixing.MixBatchHeader.parseDelimitedFrom(in);
        int n = 1 << header.getLogN();
        int layers = header.getLayers();

        Crypto.RerandomizableEncryptedMessage[][] encryptedMessages = new Crypto.RerandomizableEncryptedMessage[layers + 1][n];
        Mixing.Mix2Proof[][] proofs = null;
        if (layers > 0)
            proofs = new Mixing.Mix2Proof[header.getLayers()][n / 2];

        for (int i = 0; i < encryptedMessages.length; ++i) {
            for (int j = 0; j < encryptedMessages[i].length; ++j) {
                encryptedMessages[i][j] = Crypto.RerandomizableEncryptedMessage.parseDelimitedFrom(in);
            }
        }
        if (layers > 0) {
            for (int i = 0; i < proofs.length; ++i) {
                for (int j = 0; j < proofs[i].length; ++j) {
                    proofs[i][j] = Mixing.Mix2Proof.parseDelimitedFrom(in);
                }
            }
        }
        return new MixerOutput(header.getLogN(), proofs, encryptedMessages);
    }

    /**
     * Write to an outputstream.
     * This format can be parsed by {@link #parseDelimitedFrom(InputStream)}
     * @param out
     * @throws IOException
     */
    public void writeDelimitedTo(OutputStream out) throws IOException {
        Mixing.MixBatchHeader header = Mixing.MixBatchHeader.newBuilder()
                .setLogN(getLogN())
                .setLayers(getNumLayers())
                .build();
        header.writeDelimitedTo(out);

        for (int i = 0; i < encryptedMessages.length; ++i) {
            for (int j = 0; j < encryptedMessages[i].length; ++j) {
                encryptedMessages[i][j].writeDelimitedTo(out);
            }
        }
        if (proofs != null) {
            for (int i = 0; i < proofs.length; ++i) {
                for (int j = 0; j < proofs[i].length; ++j) {
                    proofs[i][j].writeDelimitedTo(out);
                }
            }
        }
    }

}
