extern "C" {
#include <event2/listener.h>
#include <event2/bufferevent.h>
#include <event2/buffer.h>
}

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

#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <fstream>
#include <vector>
#include <set>
#include <map>
#include <stdlib.h>
#include <string.h>

#include "evutils.h"
#include "subproblem.h"

NTL_CLIENT

struct SubproblemProgress;

typedef std::set<struct bufferevent *> BESet;
typedef std::map<struct bufferevent *, SubproblemProgress *> BEMap;

void besetdump(const BESet &bes, ostream &os)
{
    BESet::const_iterator besit;

    os << hex << "    ";
    for (besit = bes.begin(); besit != bes.end(); ++besit) {
	os << *besit << " ";
    }
    os << dec << "\n";
}

void bemapdump(const BEMap &bem, ostream &os)
{
    BEMap::const_iterator bemit;

    os << hex << "    ";
    for (bemit = bem.begin(); bemit != bem.end(); ++bemit) {
	os << bemit->first << "->" << bemit->second << " ";
    }
    os << dec << "\n";
}

struct Statuses {
    BESet idle;
    BEMap working;

    // Dump the state for debug purposes
    void dump(ostream &os) const {
	os << "  idle (" << idle.size() << "):\n";
	besetdump(idle, os);
	os << "  working (" << working.size() << "):\n";
	bemapdump(working, os);
    }
};

struct FactorDecomp {
    ZZ factor;
    vec_ZZ fvec;
};

void vsppdump(const vector<SubproblemProgress> &spv, ostream &os);

static struct ControllerState {
    ZZ rho;
    FactorDecomp p, q;
    int working;
    vector<SubproblemProgress> subproblems_p, subproblems_q;
    Statuses dpnodes, workers;

    ControllerState() : working(0) {}

    // Dump the state for debug purposes
    void dump(ostream &os) const {
	if (!working) {
	    os << "Not working\n";
	    return;
	}
	os << "P:\n";
	vsppdump(subproblems_p, os);
	os << "Q:\n";
	vsppdump(subproblems_q, os);
	os << "dpnodes:\n";
	dpnodes.dump(os);
	os << "workers:\n";
	workers.dump(os);
    }
} ctrlstate;

struct IPPort {
    unsigned char ipport[6];

    IPPort(unsigned char *ipp) {
	memmove(ipport, ipp, 6);
    }

    void dump(ostream &os) const {
	os << int(ipport[0]) << "." << int(ipport[1]) << "." <<
	    int(ipport[2]) << "." << int(ipport[3]) << ":" <<
	    ((ipport[4] << 8) + ipport[5]) << " ";
    }
};

typedef vector<IPPort> IPPortSet;

void ipportsetdump(const IPPortSet &ipps, ostream &os)
{
    IPPortSet::const_iterator ippsit;

    os << "    ";
    for (ippsit = ipps.begin(); ippsit != ipps.end(); ++ippsit) {
	ippsit->dump(os);
    }
    os << "\n";
}

struct SubproblemProgress : Subproblem {

    // The sets of dpnodes and workers currently working on this subproblem
    BESet dpnodes, workers;
    // The dpnode IPPorts registered for this subproblem
    IPPortSet ipports;

    // The desired number of DPnodes for this subproblem
    unsigned short desired_dpnodes;
    // The maximum number of workers useful for this subproblem
    unsigned int max_workers;

    // Have we found a solution?
    int solved;
    // The solution, if found.
    ZZ solution;

    SubproblemProgress(unsigned short id, const ZZ &b, const ZZ &t,
	    const ZZ &m, const ZZ &o, unsigned int dpf) :
	    Subproblem(id, b, t, m, o, dpf), solved(0) {
	// How many DPnodes should we use for a problem of this size?
	desired_dpnodes = 2;
	// How many workers would we like to use?
	ZZ sorder = SqrRoot(order >> 46);
	if (NumBits(sorder) > 30) {
	    // Just use all the workers we can find
	    max_workers = 4294967295U;  // 2^32 - 1
	} else {
	    max_workers = trunc_long(sorder,31) + 1;
	}
    }

    // Stop all dpnodes and workers and reset to unstarted state
    void reset(void) {
	BESet::iterator iter;
	unsigned char stopcmd[1] = { 'S' };

	for (BESet::iterator iter = dpnodes.begin(); iter != dpnodes.end();
		++iter) {
	    bufferevent_write(*iter, stopcmd, 1);
	    ctrlstate.dpnodes.working.erase(*iter);
	    ctrlstate.dpnodes.idle.insert(*iter);
	}
	for (BESet::iterator iter = workers.begin(); iter != workers.end();
		++iter) {
	    bufferevent_write(*iter, stopcmd, 1);
	    ctrlstate.workers.working.erase(*iter);
	    ctrlstate.workers.idle.insert(*iter);
	}
	dpnodes.clear();
	workers.clear();
	ipports.clear();
    }

    // Dump for debugging purposes
    void dump(ostream &os) const {
	os << "    dpnodes (" << dpnodes.size() << "):\n";
	besetdump(dpnodes, os);
	os << "    workers (" << workers.size() << "):\n";
	besetdump(workers, os);
	os << "    ipports (" << ipports.size() << "):\n";
	ipportsetdump(ipports, os);
    }

    void worker_write(struct bufferevent *bev) {
	bev_write(bev);
	unsigned short num_ipports = ipports.size();
	bufferevent_write(bev, &num_ipports, 2);
	for (unsigned short i = 0; i < num_ipports; ++i) {
	    bufferevent_write(bev, ipports[i].ipport, 6);
	}
    }
};

// Dump the state for debug purposes
void vsppdump(const vector<SubproblemProgress> &spv, ostream &os)
{
    vector<SubproblemProgress>::const_iterator spiter;
    int count = 0;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	++count;
	os << "  " << count << ":\n";
	spiter->dump(os);
    }
    os << "\n";
}

// Find a subproblem in the given vector that could use another DPnode,
// and give it one of the idle ones.  Only allocate it to a subproblem
// with no current DPnodes if consider_empty is true.
static void find_subproblem_for_dpnode(vector<SubproblemProgress> &spv,
	bool consider_empty)
{
    vector<SubproblemProgress>::iterator spiter;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	if (spiter->solved) continue;
	if (spiter->dpnodes.size() == 0 && consider_empty == false) continue;
	// How many DPnodes would we like to have for this subproblem?
	while (spiter->dpnodes.size() < spiter->desired_dpnodes &&
		ctrlstate.dpnodes.idle.size() > 0) {
	    // Get the first idle DPnode
	    BESet::iterator beviter = ctrlstate.dpnodes.idle.begin();

	    // Allocate it to the subproblem
	    spiter->dpnodes.insert(*beviter);
	    ctrlstate.dpnodes.working[*beviter] = &(*spiter);
	    ctrlstate.dpnodes.idle.erase(*beviter);

	    // Tell it to start listening for DPs
	    spiter->bev_write(*beviter);
	}
    }
}

// Find a subproblem in the given vector that has all of its DPnodes and
// could use another worker, and give it one of the idle ones.
static void find_subproblem_for_worker(vector<SubproblemProgress> &spv)
{
    vector<SubproblemProgress>::iterator spiter;

    for (spiter = spv.begin(); spiter != spv.end(); ++spiter) {
	if (spiter->solved) continue;

	while (spiter->ipports.size() == spiter->desired_dpnodes &&
		spiter->workers.size() < spiter->max_workers &&
		ctrlstate.workers.idle.size() > 0) {
	    // Get the first idle worker
	    BESet::iterator beviter = ctrlstate.workers.idle.begin();

	    // Allocate it to the subproblem
	    spiter->workers.insert(*beviter);
	    ctrlstate.workers.working[*beviter] = &(*spiter);
	    ctrlstate.workers.idle.erase(*beviter);

	    // Tell it to start working on the subproblem
	    spiter->worker_write(*beviter);
	}
    }
}

// See if there are any idle DPnodes or workers we can put to use
void schedule(void)
{
    cerr << "Before schedule:\n"; ctrlstate.dump(cerr);

    // Check the DPnodes

    // Iterate through the subproblems, looking for one that can use
    // another DPnode.  First look for subproblems that already have
    // some, but not all, of their DPnodes
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_p, false);
    }
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_q, false);
    }
    // If there are still more dpnodes to place, start assigning them to
    // subproblems with no current dpnodes
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_p, true);
    }
    if (ctrlstate.dpnodes.idle.size() > 0) {
	find_subproblem_for_dpnode(ctrlstate.subproblems_q, true);
    }

    // Check the workers

    // Iterate through the subproblems, looking for one that can use
    // another worker.
    if (ctrlstate.workers.idle.size() > 0) {
	find_subproblem_for_worker(ctrlstate.subproblems_p);
    }
    if (ctrlstate.workers.idle.size() > 0) {
	find_subproblem_for_worker(ctrlstate.subproblems_q);
    }

    cerr << "After schedule:\n"; ctrlstate.dump(cerr);
}

typedef enum {
    CCSTATE_START,
    CCSTATE_DPWAITRESP,
    CCSTATE_DPLISTENING,
    CCSTATE_END
} CCState;

struct ControllerConnInfo {
    CCState state;

    ControllerConnInfo() : state(CCSTATE_DPWAITRESP) {}
};

static void controller_dpnode_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    ControllerConnInfo *info = (ControllerConnInfo *)ctx;
    unsigned char cmd[1];

    while(1) {
	size_t len = evbuffer_get_length(input);
	switch (info->state) {
	    case CCSTATE_START:
	    case CCSTATE_DPWAITRESP:
		    if (len < 1) return;
		    bufferevent_read(bev, cmd, 1);
		    switch (cmd[0]) {
			case 'L':
			    info->state = CCSTATE_DPLISTENING;
			    break;
			default:
			    /* Unknown DPnode command received */
			    fprintf(stderr, "Unknown command in "
				    "controller_dpnode_reader: "
				    "%c\n", cmd[0]);
			    info->state = CCSTATE_END;
			    break;
		    }
		    break;

	    case CCSTATE_DPLISTENING:
		// Read 6 bytes
		if (len < 6) return;
		unsigned char ipport[6];
		unsigned int DPip;
		unsigned short DPport;
		bufferevent_read(bev, ipport, 6);
		memmove(&DPip, ipport, 4);
		memmove(&DPport, ipport+4, 2);
		{
		    struct in_addr DPaddr = { DPip };
		    printf("DP node at %s:%d\n", inet_ntoa(DPaddr), ntohs(DPport));
		    if (ctrlstate.dpnodes.working.count(bev) > 0) {
			ctrlstate.dpnodes.working[bev]->ipports.push_back(
			    IPPort(ipport));
			schedule();
		    }
		}
		info->state = CCSTATE_DPWAITRESP;
		break;

	    case CCSTATE_END:
		// Shut down the connection
		delete info;
		bufferevent_free(bev);
		return;
	}
    }

}

static void controller_worker_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    size_t len = evbuffer_get_length(input);

}

static void controller_dpnode_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	fprintf(stderr, "Closing dpnode connection\n");
	if (ctrlstate.dpnodes.working.count(bev)) {
	    // If we lose a dpnode from an active computation, the
	    // computation is useless.
	    SubproblemProgress *spp = ctrlstate.dpnodes.working[bev];
	    ctrlstate.dpnodes.working.erase(bev);
	    spp->dpnodes.erase(bev);
	    spp->reset();
	} else {
	    ctrlstate.dpnodes.idle.erase(bev);
	}
	delete info;
	bufferevent_free(bev);
	schedule();
    }
}

static void controller_worker_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	fprintf(stderr, "Closing worker connection\n");
	if (ctrlstate.workers.working.count(bev)) {
	    SubproblemProgress *spp = ctrlstate.workers.working[bev];
	    ctrlstate.workers.working.erase(bev);
	    spp->workers.erase(bev);
	} else {
	    ctrlstate.workers.idle.erase(bev);
	}
	delete info;
	bufferevent_free(bev);
	schedule();
    }
}

static void controller_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	fprintf(stderr, "Closing connection\n");
	ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	delete info;
	bufferevent_free(bev);
    }
}

// We're just going to read a single byte that will tell us whether the
// peer is a DPnode or a Worker
static void controller_master_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);

    size_t len = evbuffer_get_length(input);

    if (len < 1) return;

    char indata[1];
    bufferevent_read(bev, indata, 1);
    switch(indata[0]) {
	case 'D':
	    printf("DPnode\n");
	    /* Add this DPnode to the list of available ones */
	    ctrlstate.dpnodes.idle.insert(bev);
	    bufferevent_setcb(bev, controller_dpnode_reader, NULL,
		    controller_dpnode_event_cb, ctx);
	    controller_dpnode_reader(bev, ctx);
	    schedule();
	    return;
	case 'W':
	    printf("Worker\n");
	    ctrlstate.workers.idle.insert(bev);
	    bufferevent_setcb(bev, controller_worker_reader, NULL,
		    controller_worker_event_cb, ctx);
	    controller_worker_reader(bev, ctx);
	    schedule();
	    return;
	default:
	    fprintf(stderr, "Unknown command in controller_master_reader: "
		    "%c\n", indata[0]);

	    ControllerConnInfo *info = (ControllerConnInfo*)ctx;
	    delete info;
	    bufferevent_free(bev);
	    return;
    }
}

static void controller_accept_cb(struct evconnlistener *listener,
    evutil_socket_t fd, struct sockaddr *address, int socklen,
    void *ctx)
{
    // Create the state of the new connection
    ControllerConnInfo *info = new ControllerConnInfo();

    // Create a bufferevent for the new connection
    struct event_base *base = evconnlistener_get_base(listener);
    struct bufferevent *bev = bufferevent_socket_new(
	    base, fd, BEV_OPT_CLOSE_ON_FREE);

    bufferevent_setcb(bev, controller_master_reader, NULL,
	    controller_event_cb, info);

    bufferevent_enable(bev, EV_READ|EV_WRITE);
}

// Create a new controller socket.  bindport is the port to bind to (in
// host byte order), or 0 if any port will do.  ip and boundport are set
// to the IP and port of the socket, in network byte order.
void *controller_create(struct event_base *evbase, unsigned short bindport,
    unsigned int *ip, unsigned short *boundport)
{
    return listener_create(evbase, bindport, controller_accept_cb, NULL,
	ip, boundport);
}

static unsigned short curproblemid = 0;

// Take base and target mod f.factor, then decompose that into small
// subproblems given our knowledge of the factors of phi(f.factor)
static vector<SubproblemProgress> decomp(const ZZ_p &base, const ZZ_p &target,
    const FactorDecomp &f)
{
    vector<SubproblemProgress> ret;

    // Compute phi(factor)
    const int fveclen = f.fvec.length();
    ZZ phi = to_ZZ(2);
    for (int i = 0; i < fveclen; ++i) {
	phi *= f.fvec[i];
    }

    ZZ_p::init(f.factor);
    for (int i = 0; i < fveclen; ++i) {
	const ZZ& order = f.fvec[i];
	ZZ quotient = phi / order;
	ZZ_p subgroup_base = to_ZZ_p(rep(base));
	subgroup_base = power(subgroup_base, quotient);
	ZZ_p subgroup_target = to_ZZ_p(rep(target));
	subgroup_target = power(subgroup_target, quotient);

	if (subgroup_base == 1) {
	    // The original base wasn't a generator of the whole group
	    if (subgroup_target == 1) {
		// But the target is in the subgroup.  Lucky us.
		continue;
	    } else {
		ret.clear();
		return ret;
	    }
	}

	// By default, 1 in 1000 points are distinguihed points.  The
	// number in the next line is 2^32/1000
	unsigned int dpfreq = 4294967;
	if (order < 1000) {
	    // Just make every point a DP
	    dpfreq = 4294967295U;
	} else if (NumBits(order) < 27) {
	    // The frequency of DPs should be 10/sqrt(order) to avoid
	    // a DP-free cycle, so dpfreq = (10*2^32)/sqrt(order)
	    ZZ f = (to_ZZ(10) << 32) / SqrRoot(order);
	    dpfreq = trunc_long(f, 31);
	}
	ret.push_back(SubproblemProgress(curproblemid++, rep(subgroup_base),
				    rep(subgroup_target),
				    f.factor, order, dpfreq));
    }

    return ret;
}

static int generate_problem(struct event_base *evbase)
{
    // If there's already a problem on the go, don't generate another one
    if (ctrlstate.working == 1) {
	return -1;
    }
    ctrlstate.working = 1;

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

    // Decompose it mod p and mod q
    ctrlstate.subproblems_p = decomp(base, target, ctrlstate.p);
    ctrlstate.subproblems_q = decomp(base, target, ctrlstate.q);

    schedule();

    return 0;
}

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

    // Read the modulus and the factorization of its totient from cin
    cin >> ctrlstate.rho >> ctrlstate.p.factor >> ctrlstate.p.fvec >>
	ctrlstate.q.factor >> ctrlstate.q.fvec;

    unsigned short bindport = 0;

    if (argc > 1) {
	bindport = strtoul(argv[1], NULL, 10);
    }

    struct event_base *evbase = event_base_new();

    unsigned int myip;
    unsigned short myport;

    controller_create(evbase, bindport, &myip, &myport);
    struct in_addr myaddr = { myip };
    printf("Bound to %s:%d\n", inet_ntoa(myaddr), ntohs(myport));

    // Kick off the first problem to solve
    generate_problem(evbase);

    event_base_dispatch(evbase);

    return 0;
}