/*
 *  cudadl version 0.9: Compute discrete logs in smooth group orders
 *                      using CUDA
 *  Copyright (C) 2012 by Ryan Henry and Ian Goldberg
 *                        {rhenry,iang}@cs.uwaterloo.ca
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of version 3 of the GNU General Public License as
 *  published by the Free Software Foundation.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <fstream>
#include <NTL/vec_ZZ.h>
#include <NTL/ZZ_p.h>

NTL_CLIENT

// Make one attempt at generating a "smooth" prime p; i.e. a prime p
// such that all of the prime factors of p-1 are at most l_B bits long.
// Return 1 on success (and p is set to the prime and fvec is set to the
// list of factors of (p-1)/2).  Return 0 on failure.  It is expected that
// this function will fail most of the time.  p will end up being
// p_bits bits long.
static int try_gen_smooth_prime(ZZ& p, vec_ZZ &fvec, long p_bits, long l_B)
{
    ZZ prod_so_far;
    prod_so_far = 1;
    fvec.SetLength(0);

    while(NumBits(prod_so_far) < p_bits - 1) {
	long next_bits = p_bits - 1 - NumBits(prod_so_far);
	if (next_bits > l_B) next_bits = l_B;
	// For the last prime, add a bit so we're not accidentally short
	// one bit
	if (next_bits < l_B) ++next_bits;
	// There's only one prime of length < 3
	if (next_bits < 3) next_bits = 3;

	// cout << "next_bits = " << next_bits << "\n";
	ZZ nextprime = GenPrime_ZZ(next_bits);
	append(fvec, nextprime);
	prod_so_far *= nextprime;
	// cout << "next prime = " << nextprime << "\n";
    }

    // Now see if 2*prod_so_far - 1 is prime
    p = prod_so_far * 2;
    p += 1;
    return NumBits(p) == p_bits && ProbPrime(p);
}

// Verify that all of the elements of v2 are coprime to (a) all of the
// elements of v1, and (b) each other.  Return 1 if so, 0 if not.
static int all_coprime(const vec_ZZ &v1, const vec_ZZ &v2)
{
    ZZ prod_so_far;
    prod_so_far = 1;
    const int v1l = v1.length();
    for (int i=0; i < v1l; ++i) {
	prod_so_far *= v1[i];
    }

    const int v2l = v2.length();
    for (int i=0; i < v2l; ++i) {
	ZZ nextnum = v2[i];
	ZZ g = GCD(prod_so_far, nextnum);
	if (g > 1) {
	    return 0;
	}
	prod_so_far *= nextnum;
    }
    return 1;
}

// Generate rho, a product of two l_B-smooth primes.  rho should be
// about rho_bits long.  p and q will be assigned the factors of rho.
// pfvec and qfvec will be populated with the prime factors
// of (p-1)/2 and (q-1)/2 respectively.  If strong==1, then 4*rho+1 must
// also be prime.  [p and q will be 2 mod 3, so 2*p*q+1 will be
// divisible by 3 always.]
static void gen_rho(ZZ &rho, vec_ZZ &pfvec, vec_ZZ &qfvec, ZZ &p, ZZ &q,
	long rho_bits, long l_B, int strong)
{
    do {
	pfvec.SetLength(0);
	qfvec.SetLength(0);
	while(1) {
	    cerr << ".";
	    cerr.flush();
	    int ret = try_gen_smooth_prime(p, pfvec, rho_bits / 2, l_B);
	    if (ret == 1 && all_coprime(qfvec /* qfvec is empty */, pfvec)) break;
	}
	cerr << "\np = " << p << "\n\n";
	while(1) {
	    cerr << ".";
	    cerr.flush();
	    int ret = try_gen_smooth_prime(q, qfvec, rho_bits - NumBits(p), l_B);
	    if (ret == 1 && all_coprime(pfvec, qfvec)) break;
	}
	cerr << "\nq = " << q << "\n\n";
	rho = p * q;
	if (!strong) break;
	ZZ R = 4*rho + 1;
	if (!ProbPrime(R)) {
	    cerr << "4*rho + 1  = " << R << " not prime.\n\n";
	} else {
	    cerr << "R = " << R << "\n\n";
	    break;
	}
    } while (1);
}


int main(int argc, char **argv)
{
    long rho_bits = argc > 1 ? atoi(argv[1]) : 1536;
    long l_B = argc > 2 ? atoi(argv[2]) : 30;
    int strong = argc > 3 ? atoi(argv[3]) : 0;

    // Initialize the prng with some randomness from the kernel
    unsigned char randbuf[1024];
    ifstream urand("/dev/urandom");
    urand.read((char *)randbuf, sizeof(randbuf));
    urand.close();
    ZZ randzz = ZZFromBytes(randbuf, sizeof(randbuf));
    SetSeed(randzz);

    ZZ rho, p, q;
    vec_ZZ pfvec, qfvec;

    gen_rho(rho, pfvec, qfvec, p, q, rho_bits, l_B, strong);

    cout << rho << "\n" << p << "\n" << pfvec << "\n" << q << "\n"
	    << qfvec << "\n";

    return 0;
}