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

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

#include <set>
#include <map>
#ifdef LOG_MEM
#include <stdio.h>
#endif
#include <stdlib.h>
#include <string.h>
#include <errno.h>

#include <sstream>
#include <fstream>

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

// #undef VERBOSE

#ifdef SAVE_DPS
static ofstream dp_file_stream;
#endif

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;
    unsigned int id;

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

#ifdef LOG_MEM

/* Log the amount of memory used by the process.  It turns out each DP
   consumes about 338 bytes of heap.  Also, memory is not deallocated when
   the table is cleared. */
static void logmemusage(const char *label)
{
    FILE *smaps;
    FILE *logoutput;

    char logfilename[256];

    sprintf(logfilename, "mem-%05d", getpid());

    smaps = fopen("/proc/self/smaps", "r");
    if (!smaps) return;
    logoutput = fopen(logfilename, "a");
    if (!logoutput) {
	fclose(smaps);
	return;
    }
    char *line = NULL;
    size_t n = 0;
    while(getline(&line, &n, smaps) > -1) {
	if (!strstr(line, "[heap]")) {
	    free(line);
	    line = NULL;
	    continue;
	}
	// Read and parse the next line
	if (getline(&line, &n, smaps) > -1 && !strncmp(line, "Size:", 5)) {
	    unsigned long usage;
	    if (sscanf(line+5, "%lu", &usage) == 1) {
		fprintf(logoutput, "%s %lu %lu %lu\n", label,
		    dpctrlstate.table.size(), dpctrlstate.table.max_size(),
		    usage);
		break;
	    }
	}
    }
    fclose(logoutput);
    fclose(smaps);
}
#endif

static void dpnode_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
#ifdef VERBOSE
	cerr << "Closing connection " << bev << " ";
	if (events & BEV_EVENT_EOF) {
	    cerr << "EOF";
	}
	if (events & BEV_EVENT_ERROR) {
	    cerr << "ERR (" << evutil_socket_error_to_string(EVUTIL_SOCKET_ERROR()) << ")";
	}
	cerr << "\n";
#endif
	DPNodeConnInfo *info = (DPNodeConnInfo*)ctx;
	delete info;
	dpctrlstate.workers.erase(bev);
	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));
	if (dpctrlstate.current_problem == NULL) continue;
	// 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);

#ifdef SAVE_DPS
	ZZ zz_x;
	ZZFromBytes(zz_x, (const unsigned char *)(dp),
		    WORDS*sizeof(unsigned int));
	dp_file_stream << zz_x << "\n";
#endif

	pair<DTable::iterator, bool> res =
		dpctrlstate.table.insert(DTable::value_type(x,ab));
#ifdef LOG_MEM
	if (dpctrlstate.table.size() % 1000 == 0) {
	    logmemusage("insert");
	}
#endif

	if (!res.second) {
	    const ZZ& order = dpctrlstate.current_problem->order;
	    // Collision!
	    ZZ adiff = res.first->second.first - zz_a;
	    ZZ bdiff = zz_b - 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);
#ifdef VERBOSE
		cerr << "Collision after " << dpctrlstate.numdps << " DPs\n";
#endif
		unsigned char exponbytes[3+3*sizeof(unsigned int)];
		exponbytes[0] = 'E';
		memmove(exponbytes+1,
			&(dpctrlstate.current_problem->problemid), 2);
		BytesFromZZ(exponbytes+3, expon, 3*sizeof(unsigned int));
		bufferevent_write(dpctrlstate.controller_bev, exponbytes,
		    3+3*sizeof(unsigned int));
		stop_problem();
		return;
	    }
	}
	++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);
#ifdef VERBOSE
    cerr << "accepted connection " << bev << "\n";
#endif

    bufferevent_setcb(bev, dpnode_reader, NULL,
	    dpnode_event_cb, info);

    bufferevent_enable(bev, EV_READ);

    dpctrlstate.workers.insert(bev);
}

// 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)
{
    struct evconnlistener *ecl = listener_create(evbase, 0,
	dpnode_accept_cb, NULL, ip, boundport, false);

#ifdef VERBOSE
    cerr << "Listening at " << ecl << "\n";
#endif

    return ecl;
}

typedef enum {
    DPCCSTATE_AWAITCMD,
    DPCCSTATE_RDID,
    DPCCSTATE_RDPROBLEM,
    DPCCSTATE_END
} DPCCState;

struct DPControllerConnInfo {
    DPCCState state;

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

static void stop_problem(void)
{
#ifdef VERBOSE
    cerr << "Stopping problem\n";
#endif

#ifdef SAVE_DPS
    dp_file_stream.close();
#endif

    if (dpctrlstate.current_problem) {
	delete dpctrlstate.current_problem;
	dpctrlstate.current_problem = NULL;
    }
    if (dpctrlstate.listener) {
#ifdef VERBOSE
	cerr << "Closing listener " << dpctrlstate.listener << "\n";
#endif
	evconnlistener_free(dpctrlstate.listener);
	dpctrlstate.listener = NULL;
    }
    std::set<struct bufferevent *>::iterator wit;
    for (wit = dpctrlstate.workers.begin(); wit != dpctrlstate.workers.end();
	    ++wit) {
#ifdef VERBOSE
	cerr << "Closing connection " << *wit << "\n";
#endif
	bufferevent_free(*wit);
    }
    dpctrlstate.workers.clear();
#ifdef LOG_MEM
    logmemusage("preclear");
#endif
    dpctrlstate.table.clear();
#ifdef LOG_MEM
    logmemusage("postclear");
#endif
    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);
#ifdef VERBOSE
    dpctrlstate.current_problem->dump(cerr);
#endif

#ifdef SAVE_DPS
    std::ostringstream oss;
    oss << "dplist_" << dpctrlstate.id << "_" << dpctrlstate.current_problem->problemid << ".out";
    dp_file_stream.open(oss.str().c_str());
    dp_file_stream << "Subproblem " << dpctrlstate.current_problem->problemid << "\n";
#endif

    // Create the DPNode server socket
    dpctrlstate.listener = dpnode_create(bufferevent_get_base(bev),
	    &myip, &myport);
#ifdef VERBOSE
    struct in_addr myaddr = { myip };
    fprintf(stderr, "Bound to %s:%d\n", inet_ntoa(myaddr), ntohs(myport));
#endif
    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);
#ifdef VERBOSE
		cerr << "Received command " << cmd[0] << "\n";
#endif
		switch(cmd[0]) {
		    case 'I':
			info->state = DPCCSTATE_RDID;
			break;
		    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_RDID:
		if (len < sizeof(unsigned int)) return;
		unsigned int dpnode_id;
		bufferevent_read(bev, &dpnode_id, sizeof(unsigned int));
		dpctrlstate.id = dpnode_id;
		info->state = DPCCSTATE_AWAITCMD;
		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));
#ifdef VERBOSE
		cerr << "END conenction " << bev << "\n";
#endif
		dpctrlstate.workers.erase(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)) {
#ifdef VERBOSE
	fprintf(stderr, "Closing connection to controller and exiting\n");
#endif
	event_base_loopbreak(bufferevent_get_base(bev));
	bufferevent_free(bev);
    }
}

int dpnode_main(const char *controller_host, unsigned short controller_port)
{
    return controller_client(controller_host, controller_port,
				controllerconn_event_cb, false);
}