package org.gnunet.voting;

import org.gnunet.util.Program;
import org.gnunet.util.getopt.Argument;
import org.gnunet.util.getopt.ArgumentAction;

import java.math.BigInteger;

/**
 * Utilities for the modified ElGamal algorithm.
 * <p/>
 * p is a large prime, and g is a high-order element (or even a generator) of Zp*
 *
 * @author Florian Dold
 */
public class VotingParameters {
    // large prime, p = 2q + 1
    public final BigInteger p;
    // large prime, so that q divides (p-1)
    public final BigInteger q;
    // generator of Gq
    public final BigInteger g;

    public final int authorityCount;
    public final int authorityThreshold;

    public VotingParameters(BigInteger p, BigInteger q, BigInteger g, int authorityCount, int authorityThreshold) {
        this.p = p;
        this.q = q;
        this.g = g;
        this.authorityCount = authorityCount;
        this.authorityThreshold = authorityThreshold;
    }

    /**
     * which generates the p and g values from the given parameters,
     * returning the ElGamalScheme object.
     * <p/>
     * Note: can take a while...
     * 
     * @param p_bitlen bit length of p, the modulus of our group
     * @param certainty
     */
    public static VotingParameters generateRandomParameters(int p_bitlen, int certainty, int authorityCount, int authorityThreshold) {
        BigInteger[] safePrimes = generateSafePrimes(p_bitlen, certainty);
        BigInteger p = safePrimes[0];
        BigInteger q = safePrimes[1];
        BigInteger g = selectSubgroupGenerator(p, q);
        if (!g.modPow(q, p).equals(BigInteger.ONE)) {
            throw new AssertionError();
        }
        if (!(g.compareTo(p) < 0)) {
            throw new AssertionError();
        }
        return new VotingParameters(p, q, g, authorityCount, authorityThreshold);
    }

    /**
     * Finds a pair of prime BigIntegers {p, q : p = 2q + 1}, where p is
     * called a safe prime and q a sophie germain prime.
     * <p/>
     * (see: Handbook of Applied Cryptography 4.86)
     *
     * @param pBitlength bitlength of the safe prime
     * @param certainty certainty that we will really generate a pair of primes,
     *                  the probability that we fail is smaller than 2^(-certainty)
     * @return a 2-element array {p,q} of primes, where p is a safe prime
     */
    private static BigInteger[] generateSafePrimes(int pBitlength, int certainty) {
        BigInteger p, q;
        int qBitlength = pBitlength - 1;

        do {
            // generate a probably prime BigInteger
            q = new BigInteger(qBitlength, certainty, CryptoUtil.random);
            // p <- 2q + 1
            p = q.shiftLeft(1).add(BigInteger.ONE);
        } while (!p.isProbablePrime(certainty));

        return new BigInteger[]{p, q};
    }


    /**
     * Returns a higher-order-element of Gq, the subgroup of Zp*,
     * with order q where alpha is a generator of Zp*
     *
     * (see Handbook of Applied Cryptography 4.81)
     */
    private static BigInteger selectSubgroupHigherOrderElement(BigInteger alpha, BigInteger p, BigInteger q) {
        return alpha.modPow(p.subtract(BigInteger.ONE).divide(q), p);
    }

    /**
     * Get the size of the cyclic group Gp used for ElGamal
     *
     * @return the size of the cyclic group Gp used for ElGamal
     */
    public BigInteger getP() {
        return p;
    }

    /**
     * Get the generator of Gq
     *
     * @return the generator of Gq
     */
    public BigInteger getG() {
        return g;
    }

    /**
     * Get the generator of Zp*
     *
     * @return the generator of Zp*
     */
    public BigInteger getQ() {
        return q;
    }

    public static BigInteger selectSubgroupGenerator(BigInteger p, BigInteger q) {
        BigInteger alpha = selectGenerator(p, q);
        return selectSubgroupHigherOrderElement(alpha, p, q);
    }

    /**
     * Find a generator of Z_p^*, where ord(Z_p^*) = 2q.
     *
     * @param p modulus of our group
     * @param q prime factor of the order of the group Z_p^*, the other factor being 2.
     * @return generator of Z_q^*
     */
    public static BigInteger selectGenerator(BigInteger p, BigInteger q) {
        BigInteger pMinusTwo = p.subtract(BigInteger.valueOf(2));
        BigInteger g;
        /*
         * (see: Handbook of Applied Cryptography 4.80)
         */
        do {
            g = CryptoUtil.createRandomInRange(BigInteger.valueOf(2), pMinusTwo);
        }
        while (g.modPow(BigInteger.valueOf(2), p).equals(BigInteger.ONE) || g.modPow(q, p).equals(BigInteger.ONE));
        return g;
    }

    public BigInteger generateZq() {
        return CryptoUtil.createRandomInRange(BigInteger.ZERO, this.q.subtract(BigInteger.ONE));
    }

    public static void main(String... args) {
        new Program() {
            @Argument(
                    shortname = "b",
                    longname = "bits",
                    action = ArgumentAction.STORE_NUMBER,
                    description = "bit length of q")
            int bitlength = 512;
            @Argument(
                    shortname = "C",
                    longname = "certainty",
                    action = ArgumentAction.STORE_NUMBER,
                    description = "certainty")
            int certainty = 2;
            @Override
            protected void run() {
                System.out.println(String.format("Generating parameters with bitlength %s and certainty %s",
                        bitlength, certainty));
                // authority count / threshold don't matter here, just fill in a valid value ...
                VotingParameters vp = VotingParameters.generateRandomParameters(bitlength, certainty, 3, 2);
                System.out.println("p: 0x0" + vp.getP().toString(16));
                System.out.println("q: 0x0" + vp.getQ().toString(16));
                System.out.println("g: 0x0" + vp.getG().toString(16));
            }
        }.startWithoutScheduler(args);
    }

}