#include <pthread.h>

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

#include <NTL/ZZ_p.h>

#include <fstream>
#include <vector>
#include <stdio.h>
#include <errno.h>
#include <signal.h>

#include "cudadl.h"

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

NTL_CLIENT

#undef VERBOSE

typedef enum {
    WRKCCSTATE_AWAITCMD,
    WRKCCSTATE_RDPROBLEM,
    WRKCCSTATE_RDDPNODES,
    WRKCCSTATE_END
} WrkCCState;

struct WrkControllerConnInfo {
    WrkCCState state;
    unsigned short num_dpnodes;

    WrkControllerConnInfo() : state(WRKCCSTATE_AWAITCMD), num_dpnodes(0) {}
};

typedef enum { WT_NOT_RUNNING, WT_RUNNING, WT_SHOULD_STOP} WTState;

static struct WrkControllerState {
    Subproblem *current_problem;
    unsigned short num_expected_dpnodes;
    vector<struct bufferevent *> dpnodes;
    unsigned short num_connected_dpnodes;
    WTState worker_thread_state;
    pthread_t worker_thread;

    WrkControllerState(): current_problem(NULL),
	worker_thread_state(WT_NOT_RUNNING) {}
} wrkctrlstate;

// ----- Below this line are the functions running in the worker thread.

#if WORDS > 1
#define DEMUXWORD 8
#else
#define DEMUXWORD 7
#endif

// 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)
{
    unsigned int demux = dpwords[DEMUXWORD];

    struct bufferevent *bev =
	wrkctrlstate.dpnodes[demux % wrkctrlstate.num_connected_dpnodes];

    bufferevent_write(bev, dpwords+1, (WORDS+6)*sizeof(unsigned int));

    // If worker_thread_state changes to WT_SHOULD_STOP, then signal to
    // stop computation by returning true.  If for some reason, it
    // becomes WT_NOT_RUNNING (which it shouldn't), stop as well.
    return wrkctrlstate.worker_thread_state != WT_RUNNING;
}

static void *worker_thread_start(void *data)
{
    ZZ_p::init(wrkctrlstate.current_problem->modulus);
    cuda_dl(to_ZZ_p(wrkctrlstate.current_problem->base),
	    to_ZZ_p(wrkctrlstate.current_problem->target),
	    wrkctrlstate.current_problem->order,
	    wrkctrlstate.current_problem->modulus,
	    wrkctrlstate.current_problem->dpfreq, NULL);

    return NULL;
}

// ----- Above this line are the functions running in the worker thread.
//       Below are the functions running in the main (communication) thread.

static void stop_working(void)
{
#ifdef VERBOSE
    cerr << "Stopping work\n";
#endif
    if (wrkctrlstate.worker_thread_state != WT_NOT_RUNNING) {
	// Tell the worker thread to stop after its next kernel launch
	wrkctrlstate.worker_thread_state = WT_SHOULD_STOP;
	pthread_join(wrkctrlstate.worker_thread, NULL);
	wrkctrlstate.worker_thread_state = WT_NOT_RUNNING;
    }

    // Close the connections to the dpnodes
    vector<struct bufferevent *>::iterator bevit;
    for (bevit = wrkctrlstate.dpnodes.begin();
	    bevit != wrkctrlstate.dpnodes.end(); ++bevit) {
#ifdef VERBOSE
	cerr << "Closing connection to " << *bevit << "\n";
#endif
	bufferevent_free(*bevit);
    }
    wrkctrlstate.dpnodes.clear();
    wrkctrlstate.num_connected_dpnodes = 0;
    wrkctrlstate.num_expected_dpnodes = 0;
    // Careful!  Subproblem uses NTL, so we must be sure we're not
    // multithreaded at this point.
    delete wrkctrlstate.current_problem;
    wrkctrlstate.current_problem = NULL;
}

static void start_working(void)
{
#ifdef VERBOSE
    cerr << "Starting work\n";
#endif
    struct timeval now;
    gettimeofday(&now, NULL);
    char timestamp[20];
    sprintf(timestamp, "%d.%06d", (int)now.tv_sec, (int)now.tv_usec);

    cout << timestamp << ":" << output_prefix << ": Subproblem " << wrkctrlstate.current_problem->problemid << "\n";
    cout.flush();
    wrkctrlstate.worker_thread_state = WT_RUNNING;
    if (pthread_create(&wrkctrlstate.worker_thread, NULL,
			worker_thread_start, NULL)) {
	wrkctrlstate.worker_thread = WT_NOT_RUNNING;
	cerr << "Could not start worker thread\n";
    }
}

static void dpconn_event_cb(struct bufferevent *bev, short events,
    void *ctx)
{
    if (events & BEV_EVENT_CONNECTED) {
	// We have successfully connected to the dpnode
#ifdef VERBOSE
	cerr << "Connection established to dpnode " << bev << "\n";
#endif
	bufferevent_enable(bev, EV_WRITE);
	++wrkctrlstate.num_connected_dpnodes;
	if (wrkctrlstate.num_connected_dpnodes ==
		wrkctrlstate.num_expected_dpnodes) {
	    start_working();
	}
    } else if (events & (BEV_EVENT_EOF|BEV_EVENT_ERROR)) {
#ifdef VERBOSE
	cerr << "Closing connection to dpnode " << 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
	stop_working();
    }
}

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

    while(1) {
	size_t len = evbuffer_get_length(input);
	switch(info->state) {
	    case WRKCCSTATE_AWAITCMD:
		if (len < 1) return;
		bufferevent_read(bev, cmd, 1);
#ifdef VERBOSE
		cerr << "Command " << cmd[0] << " received\n";
#endif
		switch(cmd[0]) {
		    case 'P':
			info->state = WRKCCSTATE_RDPROBLEM;
			break;
		    case 'S':
			stop_working();
			break;
		    default:
			/* Unknown command received */
			fprintf(stderr, "Unknown command in "
				"controllerconn_reader: %c\n", cmd[0]);
			info->state = WRKCCSTATE_END;
			break;
		}
		break;

	    case WRKCCSTATE_RDPROBLEM:
		if (len < SUBPROBLEM_DESC_LEN+2) return;
		stop_working();
		bufferevent_read(bev, subproblem, SUBPROBLEM_DESC_LEN);
		bufferevent_read(bev, &(info->num_dpnodes), 2);
		// Careful!  Subproblem uses NTL, so we must be sure
		// we're not multithreaded at this point.
		wrkctrlstate.current_problem = new Subproblem(subproblem);
		info->state = WRKCCSTATE_RDDPNODES;
		/* FALLTHROUGH */

	    case WRKCCSTATE_RDDPNODES:
		if (len < 6*(info->num_dpnodes)) return;
		wrkctrlstate.num_expected_dpnodes = info->num_dpnodes;
		{
		    unsigned short i;
		    for(i=0;i<info->num_dpnodes;++i) {
			unsigned char ipport[6];
			bufferevent_read(bev, ipport, 6);
#ifdef VERBOSE
			cerr << "Connecting to DPnode " <<
				    int(ipport[0]) << "." <<
				    int(ipport[1]) << "." <<
			            int(ipport[2]) << "." <<
				    int(ipport[3]) << ":" <<
				    ((ipport[4] << 8) + ipport[5]) <<
				    "\n";
#endif
			struct bufferevent *dpbev = client_create(
				bufferevent_get_base(bev), ipport,
				dpconn_event_cb, true);
#ifdef VERBOSE
			cerr << "Starting connection to dpnode " << dpbev << "\n";
#endif
			if (dpbev) {
			    wrkctrlstate.dpnodes.push_back(dpbev);
			} else {
			    stop_working();
			}
		    }
		}
		info->state = WRKCCSTATE_AWAITCMD;
		break;

	    case WRKCCSTATE_END:
		// Shut down
		delete info;
		event_base_loopbreak(bufferevent_get_base(bev));
#ifdef VERBOSE
		cerr << "Closing connection to " << bev << "\n";
#endif
		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] = { 'W' };

        bufferevent_enable(bev, EV_READ|EV_WRITE);
	bufferevent_write(bev, id, 1);
	bufferevent_setcb(bev, controllerconn_reader, NULL,
		controllerconn_event_cb, new WrkControllerConnInfo());
    } 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));
#ifdef VERBOSE
	cerr << "Closing connection to " << bev << "\n";
#endif
	bufferevent_free(bev);
    }
}

int worker_main(const char *controller_host, unsigned short controller_port)
{
    // 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);

    evthread_use_pthreads();

    signal(SIGPIPE, SIG_IGN);

    return controller_client(controller_host, controller_port,
				controllerconn_event_cb, true);
}