/*
 *  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 <sys/time.h>
#include <NTL/vec_ZZ.h>
#include <NTL/ZZ_p.h>

#include <map>

#include <sys/types.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <cuda_runtime.h>

#include "cudadl.h"

NTL_CLIENT

string output_prefix;

typedef map<std::string, pair<ZZ,ZZ> > DTable;

struct CBData {
    const ZZ_p &base;
    const ZZ_p &target;
    const ZZ &order;
    unsigned long long numdp;
    DTable dtable;
    bool found_collision;
    ZZ expon;

    CBData(const ZZ_p &_base, const ZZ_p &_target, const ZZ &_order) :
	base(_base), target(_target), order(_order), numdp(0),
	found_collision(false) {}
};

// This function is called from inside cuda_dl for each DP it encounters.
// It calls the function named "dpcallback" directly.  It would be
// cleaner if this were passed as a function pointer to cuda_dl, but
// that makes nvcc 3.1 segfault. :-p
// dp points to an array of WORDS+7 unsigned ints:
// - 1 word of threadID/blockID
// - WORDS words of the dp value
// - 3 words of a
// - 3 words of b
bool dpcallback(void *cbdata, unsigned int *dpwords)
{
    CBData *d = (CBData*)cbdata;
    // WARNING: this assumes
    //    sizeof(unsigned long) == sizeof(unsigned long long) !
    ZZ zz_a = to_ZZ(dpwords[WORDS+3]);
    zz_a <<= 32;
    zz_a += dpwords[WORDS+2];
    zz_a <<= 32;
    zz_a += dpwords[WORDS+1];
    ZZ zz_b = to_ZZ(dpwords[WORDS+6]);
    zz_b <<= 32;
    zz_b += dpwords[WORDS+5];
    zz_b <<= 32;
    zz_b += dpwords[WORDS+4];

    //ZZ_p dp = power(d->base, zz_a) * power(d->target, zz_b);

    // cerr << "DP " << ++(d->numdp) << "\r";
    pair<ZZ,ZZ> ab(zz_a,zz_b);
    pair<DTable::iterator, bool> res;
    string x((const char *)(dpwords+1), WORDS*sizeof(unsigned int));

    res = d->dtable.insert(DTable::value_type(x, ab));
    if (!res.second) {
	// Collision!
	ZZ adiff = to_ZZ(res.first->second.first) - zz_a;
	ZZ bdiff = zz_b - to_ZZ(res.first->second.second);
	if (bdiff < 0) bdiff += d->order;
	ZZ binv;
	if (InvModStatus(binv, bdiff, d->order) == 0) {
	    d->expon = MulMod(binv, adiff, d->order);
	    d->found_collision = true;
	} else {
	    if (d->order > (1<<20)) {
		cerr << "Unhelpful collision\n";
	    }
	}
    }

    return d->found_collision;
}

// Compute the discrete log of target mod p, to the given base.
// p must be the current ZZ_p modulus.
// Place the result in exp.  fvec is a vector of the factors of
// (p-1)/2, which must each be small enough to compute discrete logs
// with some other method (kangaroo, index calculus, GNFS, etc.).
// label is "p" or "q", to be printed to report progress.
// Return 0 on failure, 1 on success.
static int p_dl(const ZZ_p &target, const ZZ_p &base, ZZ &exp,
    const vec_ZZ &fvec, const string &label)
{
    const int flen = fvec.length();

    // Compute phi(p)
    ZZ phip;
    phip = 2;
    for (int i = 0; i < flen; ++i) {
	phip *= fvec[i];
    }

    // Invariant: the desired exp \equiv (curexp mod curmodulus)
    ZZ curmodulus, curexp;
    curmodulus = 2;
    curexp = 0;

    for (int i = 0; i < flen; ++i) {
	cout << label << " submodulus " << i+1 << " of " << flen << "... ";
	cout.flush();
	// Figure out exp mod fvec[i] by taking each side to the power
	// of phirho/fvec[i] so that we're working in the
	// subgroup of order fvec[i].
	ZZ quotient = phip / fvec[i];
	ZZ_p subgroup_base = power(base, quotient);
	ZZ_p subgroup_target = power(target, quotient);

	if (subgroup_base == 1) {
	    // The original base wasn't a generator of the whole group.
	    if (subgroup_target == 1) {
		cout << "Non-unique solution (mod " << fvec[i] <<")\n";
		continue;
	    } else {
		cout << "Target not in subgroup generated by base\n";
		return 0;
	    }
	}

	// Now use your favourite method to get the DL of
	// subgroup_target with base subgroup_base, knowing that it's in
	// the range [0,fvec[i]).

	cout << "\n";
	ZZ md = ZZ_p::modulus();
	CBData cbdata(subgroup_base, subgroup_target, fvec[i]);
	struct timeval st, et;
	unsigned int dpfreq = 4294967; // 2^32/1000
	if (fvec[i] < 1000000) {
	    dpfreq = 4294967295; // 2^32-1 : every point is a DP
	}
	gettimeofday(&st, NULL);
	cuda_dl(subgroup_base, subgroup_target, fvec[i], md, dpfreq, &cbdata);
	ZZ subgroup_dl = cbdata.expon;
	gettimeofday(&et, NULL);
	unsigned long us_elapsed = (et.tv_sec-st.tv_sec)*1000000 +
	    (et.tv_usec-st.tv_usec);
	printf("%ld.%06ld seconds elapsed\n", us_elapsed/1000000,
	    us_elapsed % 1000000);

	CRT(curexp, curmodulus, subgroup_dl, fvec[i]);
	// cout << "CRT\n";
	// cout << "curexp = " << curexp << "\n";
	// cout << "curmodulus = " << curmodulus << "\n\n";
    }
    // We'd like a non-negative answer back
    if (curexp >= 0) {
	exp = curexp;
    } else {
	exp = curexp + curmodulus;
    }
    return 1;
}

typedef struct {
    pid_t pid;
    int rfd;
} PDLHandle;

// Behave like p_dl, but do the work in an asynchronous subprocess
static PDLHandle* p_dl_fork_start(const ZZ &p, const ZZ_p &target,
    const ZZ_p &base, const vec_ZZ &fvec, const string &label, int deviceid)
{
    PDLHandle *handle = new PDLHandle;
    int fds[2];
    int res = socketpair(AF_UNIX, SOCK_STREAM, 0, fds);
    if (res < 0) {
	perror("socketpair");
	delete handle;
	return NULL;
    }
    handle->rfd = fds[0];
    pid_t childpid = fork();
    if (childpid == -1) {
	perror("fork");
	delete handle;
	return NULL;
    }
    if (childpid == 0) {
	// Child
	close(fds[0]);
	int wfd = fds[1];

	ZZ exp;
	ZZ_p::init(p);
	ZZ_p target_p, base_p;
	conv(target_p, rep(target));
	conv(base_p, rep(base));

	cudaError_t cudares = cudaSetDevice(deviceid);
	if (cudares != cudaSuccess) {
	    cerr << "Error setting CUDA device\n";
	    exit(1);
	}
	int res = p_dl(target_p, base_p, exp, fvec, label);

	if (res) {
	    // Write the result back to the parent using wfd
	    unsigned short explen = NumBytes(exp);
	    res = write(wfd, &explen, sizeof(unsigned short));
	    unsigned char expbuf[explen];
	    BytesFromZZ(expbuf, exp, explen);
	    res = write(wfd, expbuf, explen);
	    close(wfd);
	}

	exit(0);
    } else {
	// Parent
	close(fds[1]);
    }

    return handle;
}

// Wait until the subprocess started by p_dl_fork_start completes, and
// return its result.  handle is cleaned up.
static int p_dl_fork_join(PDLHandle *handle, ZZ &exp)
{
    if (!handle) return 0;
    pid_t childpid = handle->pid;
    unsigned short explen;
    int res = read(handle->rfd, &explen, sizeof(unsigned short));
    if (res < (int)sizeof(unsigned short)) {
	close(handle->rfd);
	delete handle;
	waitpid(childpid, NULL, 0);
	return 0;
    }
    unsigned char expbuf[explen];
    res = read(handle->rfd, expbuf, explen);
    if (res < explen) {
	close(handle->rfd);
	delete handle;
	waitpid(childpid, NULL, 0);
	return 0;
    }
    ZZFromBytes(exp, expbuf, explen);
    close(handle->rfd);
    delete handle;
    waitpid(childpid, NULL, 0);
    return 1;
}

int main(int argc, char **argv)
{
    // 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;
    cin >> rho >> p >> pfvec >> q >> qfvec;

    // Generate a DLP mod rho (in the large odd-order subgroup)
    ZZ_p::init(rho);
    ZZ_p base = power(random_ZZ_p(), 2);
    ZZ_p target = power(random_ZZ_p(), 2);

    cout << "base = " << base << "\n";
    cout << "target = " << target << "\n";

    ZZ exp_p, exp_q, exp;
    int res_p = 0;
    int res_q = 0;

    PDLHandle *handle_p, *handle_q;

    handle_p = p_dl_fork_start(p, target, base, pfvec, "p", 0);
    handle_q = p_dl_fork_start(q, target, base, qfvec, "q", 1);
    res_p = p_dl_fork_join(handle_p, exp_p);
    res_q = p_dl_fork_join(handle_q, exp_q);

    if (res_p && res_q) {
	ZZ_p::init(rho);
	ZZ pm1 = (p - 1)/2;
	ZZ qm1 = (q - 1)/2;
	if (exp_p < 0) exp_p += pm1;
	if (exp_q < 0) exp_q += qm1;
	CRT(exp_p, pm1, exp_q, qm1);
	exp = exp_p;
	// We'd like a non-negative answer back
	if (exp < 0) exp += pm1;

	cout << "exp = " << exp << "\n";
	ZZ_p base_exp;
	power(base_exp, base, exp);
	if (base_exp == target) {
	    cout << "CORRECT!\n";
	} else {
	    cout << "INCORRECT:\nbase^exp = " << base_exp << "\n";
	    cout << "target   = " << target << "\n";
	}
    } else {
	cout << "FAIL\n";
    }

    return 0;
}