package meerkat.crypto.dkg.feldman;

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.secretsharing.shamir.Polynomial;
import meerkat.crypto.secretsharing.shamir.SecretSharing;
import meerkat.crypto.utils.BigIntegerByteEncoder;
import meerkat.crypto.utils.GenerateRandomPrime;
import org.factcenter.qilin.primitives.Group;
import org.factcenter.qilin.primitives.concrete.Zpstar;
import org.factcenter.qilin.util.ByteEncoder;
import org.junit.Test;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Random;
import java.util.Set;

/**
 * Created by Tzlil on 3/21/2016.
 */
public class DKGTest {

    int 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 t = 9;
    int n = 20;

    public void oneTest(Testable testable) throws Exception {
        for (int i = 0; i < testable.threads.length ; i++){
            testable.threads[i].start();
        }
        for (int i = 0; i < testable.threads.length ; i++){
            testable.threads[i].join();
        }

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

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


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

        for (int i: testable.valids){
            sharesList.add(testable.dkgs[i - 1].getShare());
        }
        Polynomial.Point[] shares = new Polynomial.Point[sharesList.size()];
        for (int i = 0; i < shares.length; i ++){
            shares[i] = sharesList.get(i);
        }

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


    @Test
    public void test() throws Exception {
        Testable testable;
        for (int i = 0; i < tests; i++){
            testable = new Testable(new Random());
            oneTest(testable);
        }
    }

    class Testable{
        Set<Integer> valids;
        Set<Integer> QUAL;
        Set<Integer> aborted;
        Set<Integer> malicious;
        User<BigInteger>[] dkgs;
        Thread[] threads;
        BigInteger g;
        BigInteger secret;

        public Testable(Random random) {
            this.dkgs = new User[n];
            this.valids = new HashSet<Integer>();
            this.QUAL = new HashSet<Integer>();
            this.aborted = new HashSet<Integer>();
            this.malicious = new HashSet<Integer>();
            this.threads = new Thread[n];
            this.g = sampleGenerator(random);
            ArrayList<Integer> ids = new ArrayList<Integer>();
            for (int id = 1; id<= n ; id++){
                ids.add(id);
            }
            int id;
            BigInteger s;
            Protocol<BigInteger> dkg;
            this.secret = BigInteger.ZERO;
            ChannelImpl channels = new ChannelImpl();
            ByteEncoder<BigInteger> byteEncoder = new BigIntegerByteEncoder();
            while (!ids.isEmpty()) {
                id = ids.remove(random.nextInt(ids.size()));
                Channel channel  = channels.getChannel(id);
                s = randomIntModQ(random);
                dkg = new meerkat.crypto.dkg.feldman.Protocol<BigInteger>(t, n, s, random, q, g, group, id,byteEncoder);
                dkgs[id - 1] = randomDKGUser(id,channel,dkg,random);
                threads[id - 1] = new Thread(dkgs[id - 1]);
                if(QUAL.contains(id)){
                    this.secret =  this.secret.add(s).mod(q);
                }
            }

        }

        public User<BigInteger> randomDKGUser(int id, Channel channel, Protocol<BigInteger> dkg, Random random){
            if (QUAL.size() <= t) {
                valids.add(id);
                QUAL.add(id);
                return new User<BigInteger>(dkg,channel);
            }else{
                int type = random.nextInt(3);
                switch (type){
                    case 0:// regular
                        valids.add(id);
                        QUAL.add(id);
                        return new User<BigInteger>(dkg,channel);
                    case 1:// abort
                        int abortStage = random.nextInt(2) + 1; // 1 or 2
                        aborted.add(id);
                        if (abortStage == 2){
                            QUAL.add(id);
                        }
                        return new DKGUserImplAbort(dkg,channel,abortStage);
                    case 2:// malicious
                        malicious.add(id);
                        Set<Integer> falls = DKGMaliciousUser.selectFallsRandomly(valids,random);
                        Protocol<BigInteger> maliciousDKG = DKGMaliciousUser.generateMaliciousDKG(dkg,channel,random);
                        return new DKGMaliciousUser(dkg,maliciousDKG,channel,falls);
                    default:
                        return null;
                }
            }
        }

        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);
        }

    }
}
