package meerkat.crypto.dkg.gjkr;

import meerkat.comm.ChannelImpl;
import meerkat.crypto.utils.Arithmetic;
import meerkat.crypto.utils.concrete.Fp;
import meerkat.comm.Channel;
import meerkat.crypto.secretsharing.feldman.VerifiableSecretSharing;
import meerkat.crypto.dkg.feldman.DKGMaliciousUser;
import meerkat.crypto.secretsharing.shamir.Polynomial;
import meerkat.crypto.secretsharing.shamir.SecretSharing;
import meerkat.crypto.utils.BigIntegerByteEncoder;
import meerkat.crypto.utils.GenerateRandomPrime;
import meerkat.protobuf.Crypto;
import org.factcenter.qilin.primitives.Group;
import org.factcenter.qilin.primitives.concrete.Zpstar;
import org.factcenter.qilin.util.ByteEncoder;
import org.junit.Assert;
import org.junit.Test;
import org.junit.internal.runners.statements.Fail;

import static org.junit.Assert.*;

import java.math.BigInteger;
import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

/**
 * Created by Tzlil on 3/29/2016.
 * TODO: Separate into multiple tests,
 * TODO: Make tests deterministic (using constant seed for random generator)
 */
public class SDKGTest {

    private ExecutorService executorService = Executors.newCachedThreadPool();

    final static int NUM_TESTS = 10;
    BigInteger p = GenerateRandomPrime.SafePrime100Bits;
    BigInteger q = p.subtract(BigInteger.ONE).divide(BigInteger.valueOf(2));
    Group<BigInteger> group = new Zpstar(p);
    Arithmetic<BigInteger> arithmetic = new Fp(q);
    int initialT = 17;
    int initialN = 20;
    Random rand = new Random(1);

    public void oneTest(Testable testable) throws Exception {
        for (int i = 0; i < testable.sdkgs.length; i++){
            testable.futures[i] = executorService.submit(testable.sdkgs[i]);
        }
        for (Future ftr : testable.futures) {
            ftr.get();
        }

        // got the right public value
        BigInteger publicValue = group.multiply(testable.g, testable.secret);
        for (int i : testable.valids){
            assertEquals (testable.sdkgs[i-1].getPublicValue(), publicValue);
        }

        // assert valid verification values
        BigInteger expected,verification;
        for (int i : testable.valids){
            expected = group.multiply(testable.g, testable.sdkgs[i - 1].getShare().y);
            verification = VerifiableSecretSharing.computeVerificationValue(i, testable.sdkgs[i - 1].getCommitments(), group);
            assertEquals (expected, verification);
        }


        // restore the secret from shares
        ArrayList<Polynomial.Point> sharesList = new ArrayList<>();

        for (int i : testable.valids){
            sharesList.add(testable.sdkgs[i - 1].getShare());
        }
        Polynomial.Point[] shares = new Polynomial.Point[sharesList.size()];
        shares = sharesList.toArray(shares);

        BigInteger calculatedSecret = SecretSharing.recoverSecret(shares, arithmetic);
        assertEquals (calculatedSecret, testable.secret);
    }


    @Test
    public void runSharingProtocol() throws Exception {
        Testable testable;
        for (int i = 0; i < NUM_TESTS; i++) {
            testable = new Testable(initialN+i, initialT+i, group, q, rand);
            oneTest(testable);
        }
    }


    static class Testable {
        Set<Integer> valids;
        Set<Integer> QUAL;
        Set<Integer> aborted;
        Set<Integer> malicious;
        User<BigInteger>[] sdkgs;
        Future<?>[] futures;
        BigInteger g;
        BigInteger h;
        BigInteger secret;
        Group<BigInteger> group;
        int n;
        int t;
        BigInteger q;
        Random random;
        ChannelImpl channels = new ChannelImpl();

        public Testable(int n, int t, Group<BigInteger> group, BigInteger q, Random random) {
            this.n = n;
            this.t = t;
            this.group = group;
            this.q = q;
            this.random = random;
            this.sdkgs = new User[n];
            this.valids = new HashSet<>();
            this.QUAL = new HashSet<>();
            this.aborted = new HashSet<>();
            this.malicious = new HashSet<>();
            this.futures = new Future[n];
            this.g = sampleGenerator(random);
            this.h = group.multiply(g, randomIntModQ(random));
            List<Integer> ids = new ArrayList<>();
            for (int id = 1; id<= n ; id++){
                ids.add(id);
            }
            int id;
            BigInteger s;
            Channel channel;
            Protocol<BigInteger> sdkg;
            this.secret = BigInteger.ZERO;
            ByteEncoder<BigInteger> encoder = new BigIntegerByteEncoder();
            while (!ids.isEmpty()) {
                id = ids.remove(random.nextInt(ids.size()));
                s = randomIntModQ(random);
                channel  = channels.getChannel(id);
                sdkg = new Protocol<>(t, n, s, random, q, g , h, group, id, encoder);
                sdkgs[id - 1] = randomSDKGUser(id, channel, sdkg);
                if(QUAL.contains(id)){
                    this.secret =  this.secret.add(s).mod(q);
                }
            }

        }

        enum UserType {
            HONEST,
            FAILSTOP,
            MALICIOUS,
        }

        public User<BigInteger> newSDKGUser(int id, Channel channel, Protocol<BigInteger> sdkg, UserType userType) {
            switch(userType) {
                case HONEST:
                    valids.add(id);
                    QUAL.add(id);
                    return new User<>(sdkg,channel);

                case FAILSTOP:
                    int abortStage = random.nextInt(3) + 1; // 1 or 2 or 3
                    aborted.add(id);
                    if (abortStage > 1){
                        QUAL.add(id);
                    }
                    return new SDKGUserImplAbort(sdkg, channel, abortStage);

                case MALICIOUS:
                    malicious.add(id);
                    Set<Integer> falls = DKGMaliciousUser.selectFallsRandomly(valids, random);
                    Protocol<BigInteger> maliciousSDKG = SDKGMaliciousUserImpl.generateMaliciousSDKG(sdkg, channel, random);
                    return new SDKGMaliciousUserImpl(sdkg, maliciousSDKG, channel, falls);

            }
            fail("Unknown user type");
            return null;
        }

        public User<BigInteger> randomSDKGUser(int id, Channel channel, Protocol<BigInteger> sdkg){
            if (QUAL.size() <= t) {
                return newSDKGUser(id, channel, sdkg, UserType.HONEST);
            } else {
                UserType type = UserType.values()[random.nextInt(UserType.values().length)];
                return newSDKGUser(id, channel, sdkg, type);
            }
        }

        public BigInteger sampleGenerator(Random random){
            BigInteger ZERO = group.zero();
            BigInteger g;
            do {
                g = group.sample(random);
            } while (!g.equals(ZERO) && !group.multiply(g, q).equals(ZERO));
            return g;
        }

        public BigInteger randomIntModQ(Random random){
            return new BigInteger(q.bitLength(), random).mod(q);
        }

    }


}
