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

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

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

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

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

typedef enum {
    DPSTATE_START,
    DPSTATE_END
} DPState;

struct DPNodeConnInfo {
    DPState state;

    DPNodeConnInfo() : state(DPSTATE_START) {}
};

static struct DPControllerState {
    struct bufferevent *controller_bev;
    Subproblem *current_problem;
    struct evconnlistener *listener;
    std::set<struct bufferevent *> workers;
    DTable table;
    unsigned long long numdps;

    DPControllerState() : controller_bev(NULL), current_problem(NULL),
			    listener(NULL) {}
} dpctrlstate;

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

static void stop_problem(void);

static void dpnode_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    unsigned int dp[WORDS+6];

    while(1) {
	size_t len = evbuffer_get_length(input);
	if (len < ((WORDS+6)*sizeof(unsigned int))) break;
	bufferevent_read(bev, dp, (WORDS+6)*sizeof(unsigned int));
	// The (WORDS+6) unsigned ints we read are:
	// - WORDS words for the value of the dp
	// - 3 words for a
	// - 3 words for b
	ZZ zz_a, zz_b;
	ZZFromBytes(zz_a, (const unsigned char *)(dp+WORDS),
		3*sizeof(unsigned int));
	ZZFromBytes(zz_b, (const unsigned char *)(dp+WORDS+3),
		3*sizeof(unsigned int));
	string x((const char *)(dp), WORDS*sizeof(unsigned int));
	pair<ZZ,ZZ> ab(zz_a,zz_b);
	pair<DTable::iterator, bool> res =
		dpctrlstate.table.insert(DTable::value_type(x,ab));

	if (!res.second) {
	    const ZZ& order = dpctrlstate.current_problem->order;
	    // Collision!
	    ZZ adiff = to_ZZ(res.first->second.first) - zz_a;
	    ZZ bdiff = zz_b - to_ZZ(res.first->second.second);
	    while (bdiff < 0) bdiff += order;
	    while (bdiff >= order) bdiff -= order;
	    while (adiff < 0) adiff += order;
	    while (adiff >= order) adiff -= order;
	    ZZ binv;
	    if (InvModStatus(binv, bdiff, order) == 0) {
		ZZ expon = MulMod(binv, adiff, order);
		cerr << "Collision after " << dpctrlstate.numdps << " DPs\n";
		unsigned char exponbytes[3*sizeof(unsigned int)];
		BytesFromZZ(exponbytes, expon, 3*sizeof(unsigned int));
		bufferevent_write(dpctrlstate.controller_bev, exponbytes,
		    3*sizeof(unsigned int));
		stop_problem();
	    }
	}
	++dpctrlstate.numdps;
    }
}

static void dpnode_accept_cb(struct evconnlistener *listener,
    evutil_socket_t fd, struct sockaddr *address, int socklen,
    void *ctx)
{
    DPNodeConnInfo *info = new DPNodeConnInfo();

    // 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, dpnode_reader, NULL,
	    dpnode_event_cb, info);

    bufferevent_enable(bev, EV_READ);
}

// Create a new DPnode socket.  ip and boundport are set to the IP and
// port of the socket, in network byte order.
struct evconnlistener *dpnode_create(struct event_base *evbase,
    unsigned int *ip, unsigned short *boundport)
{
    return listener_create(evbase, 0, dpnode_accept_cb, NULL,
	ip, boundport, false);
}

typedef enum {
    DPCCSTATE_AWAITCMD,
    DPCCSTATE_RDPROBLEM,
    DPCCSTATE_END
} DPCCState;

struct DPControllerConnInfo {
    DPCCState state;

    DPControllerConnInfo() : state(DPCCSTATE_AWAITCMD) {}
};

static void stop_problem(void)
{
    if (dpctrlstate.current_problem) {
	dpctrlstate.current_problem = NULL;
    }
    if (dpctrlstate.listener) {
	evconnlistener_free(dpctrlstate.listener);
	dpctrlstate.listener = NULL;
    }
    std::set<struct bufferevent *>::iterator wit;
    for (wit = dpctrlstate.workers.begin(); wit != dpctrlstate.workers.end();
	    ++wit) {
	bufferevent_free(*wit);
    }
    dpctrlstate.workers.clear();
    dpctrlstate.table.clear();
    dpctrlstate.numdps = 0;
}

static void start_problem(struct bufferevent *bev,
	const unsigned char *subproblem)
{
    unsigned int myip;
    unsigned short myport;

    stop_problem();
    dpctrlstate.current_problem = new Subproblem(subproblem);
    dpctrlstate.current_problem->dump(cerr);

    // Create the DPNode server socket
    dpctrlstate.listener = dpnode_create(bufferevent_get_base(bev),
	    &myip, &myport);
    struct in_addr myaddr = { myip };
    fprintf(stderr, "Bound to %s:%d\n", inet_ntoa(myaddr), ntohs(myport));
    unsigned char idstring[7];
    idstring[0] = 'L';
    memmove(idstring+1, &myip, 4);
    memmove(idstring+5, &myport, 2);

    bufferevent_write(bev, idstring, 7);
}

static void controllerconn_reader(struct bufferevent *bev, void *ctx)
{
    struct evbuffer *input = bufferevent_get_input(bev);
    DPControllerConnInfo *info = (DPControllerConnInfo *)ctx;
    unsigned char cmd[1];
    unsigned char subproblem[SUBPROBLEM_DESC_LEN];

    while(1) {
	size_t len = evbuffer_get_length(input);
	switch(info->state) {
	    case DPCCSTATE_AWAITCMD:
		if (len < 1) return;
		bufferevent_read(bev, cmd, 1);
		switch(cmd[0]) {
		    case 'P':
			info->state = DPCCSTATE_RDPROBLEM;
			break;
		    case 'S':
			stop_problem();
			break;
		    default:
			/* Unknown command received */
			fprintf(stderr, "Unknown command in "
				"controllerconn_reader: %c\n", cmd[0]);
			info->state = DPCCSTATE_END;
			break;
		}
		break;

	    case DPCCSTATE_RDPROBLEM:
		if (len < SUBPROBLEM_DESC_LEN) return;
		bufferevent_read(bev, subproblem, SUBPROBLEM_DESC_LEN);
		start_problem(bev, subproblem);
		info->state = DPCCSTATE_AWAITCMD;
		break;

	    case DPCCSTATE_END:
		// Shut down
		delete info;
		event_base_loopbreak(bufferevent_get_base(bev));
		bufferevent_free(bev);
		return;
	}
    }
}

static void controllerconn_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & BEV_EVENT_CONNECTED) {
	// We have successfully connected to the controller
	char id[1] = { 'D' };

        bufferevent_enable(bev, EV_READ|EV_WRITE);
	bufferevent_write(bev, id, 1);
	bufferevent_setcb(bev, controllerconn_reader, NULL,
		controllerconn_event_cb, new DPControllerConnInfo());
	dpctrlstate.controller_bev = bev;
    } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
	fprintf(stderr, "Closing connection to controller and exiting\n");
	event_base_loopbreak(bufferevent_get_base(bev));
	bufferevent_free(bev);
    }
}

int main(int argc, char **argv)
{
    if (argc != 3) {
	fprintf(stderr, "Usage: %s controller_host controller_port\n", argv[0]);
	return 1;
    }

    unsigned short controller_port = strtoul(argv[2], NULL, 10);

    return controller_client(argv[1], controller_port,
				controllerconn_event_cb, false);
}