#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <pm.h>
#include <score.h>


enum operation {
	NOP, PINGPONG, REPLY, BURST, SINK,
	VREPLY, VWRITE, VREAD, 
};

#define	BUFSIZE		(128*1024*1024 + 8192)	/* must be larger than cache */

#define ITERATION	100000			/* default iteration	*/
#define MSG_LENGTH	8			/* default message length */
#define TIME_QUANTUM	100000			/* default context switch */
#define TIME_LOOP	  1000			/* default wait time for loop */

pmContext	*pmc;
int		nprocs;
int		rank;
size_t		*mtu;
char		*data;

int	bw_flag = 0;

int	pe = 1,
	self = 0,
	channel = 0,
	context = 0,
	dest = -1,
	length = MSG_LENGTH,
	iteration = ITERATION,
	time_quantum = TIME_QUANTUM,
	second = 0,
	cache = 4 * 1024 * 1024,
	desc = -1,
	debug = 0;

void
myexception(int sig, int code, struct sigcontext *sc, char* addr)
{
#ifdef i386
  struct sigcontext *sbx = (struct sigcontext *)&code;
  score_dump_registers(sig, code, sbx);
#endif /* i386  */
#ifdef __alpha__
  score_dump_registers(sig, code, sc);
#endif
#ifdef __ia64__
  score_dump_registers(sig, code, sc);
#endif
  score_attach_debugger("Exception signal", sig);
  fclose(stdout);
  fclose(stderr);
  _exit(sig);                 /* do not call exit() */
  /* do not reach here */
}

void
sig_handler(int sig)
{

	pmDumpContext(pmc, stderr);
	if (sig == SIGINT)
		exit(0);
}

void
fatal(char *msg, int error)
{

	pmDumpContext(pmc, stderr);
	exit(1);
}

struct timeval start_tv, end_tv;

#define	TV2DOUBLE(t)	((t)->tv_sec + (t)->tv_usec / 1000000.0)

void
print_time(size_t length, int iteration,
	   struct timeval *startp, struct timeval *endp)
     /*	   double st, double et)*/
{
	double	delta;

	delta = TV2DOUBLE(endp) - TV2DOUBLE(startp);
	/*	delta = et - st;*/
	if (bw_flag)
		printf("%Zd\t%g\n", length, length / (delta / iteration));
	else
		printf("%Zd\t%g\n", length,  delta / iteration);
}


void
scoreinit(int argc, char **argv)
{
    int		i;

    score_initialize(argc, argv);
    nprocs= score_num_node;
    rank = score_self_node;
    signal(SIGHUP,  (void(*)(int))myexception);
    signal(SIGILL,  (void(*)(int))myexception);
    signal(SIGFPE,  (void(*)(int))myexception);
    signal(SIGBUS,  (void(*)(int))myexception);
    signal(SIGSEGV, (void(*)(int))myexception);

    pmc = score_pmnet[0];
}

int score_runtime_resource( void ) { return 0; }

void
pingpong(pmContext *pmc, int dest, size_t length, int n)
{
	int	i, error;
	long long k;
	caddr_t	addr, addr2 = NULL;
	size_t	length2, mtu;

	for (i = 0; i < n; i++) {
		//printf("[%d] ping to %d[%d]\n", rank, dest, i);
		//fflush(stdout);

		while ((error = pmGetSendBuffer(pmc, dest, &addr, length)) == ENOBUFS)
			;
		if (error != PM_SUCCESS)
			fatal("pmGetSendBuffer", error);

		*(int *)addr = htonl(self);
		if ((error = pmSend(pmc)) != PM_SUCCESS)
			fatal("pmSend", error);
		if (addr2 != NULL) {
			if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
				fatal("pmReleaseReceiveBuffer", error);
		}
		while ((error = pmReceive(pmc, &addr2, &length2)) == ENOBUFS)
			;
		if (error != PM_SUCCESS) {
		  	fprintf(stderr,"error loop %d\n", i);
			fatal("pmReceive", error);
		}
	}
	if (addr2 != NULL) {
		if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
			fatal("pmReleaseReceiveBuffer", error);
	}
}

void
reply(pmContext *pmc)
{
	int	error, dest;
	caddr_t	addr, addr2;
	size_t	length;

	for (;;) {
		while ((error = pmReceive(pmc, &addr, &length)) == ENOBUFS)
			;
		if (error != PM_SUCCESS)
			fatal("pmReceive", error);
		dest = ntohl(*(int *)addr);

		//printf("[%d] recv from %d\n", rank, dest);
		//fflush(stdout);

		while ((error = pmGetSendBuffer(pmc, dest, &addr2, length)) == ENOBUFS)
			;
		if (error != PM_SUCCESS)
			fatal("pmGetSendBuffer", error);
		*(int *)addr2 = htonl(self);
		if ((error = pmSend(pmc)) != PM_SUCCESS)
			fatal("pmSend", error);
		if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
			fatal("pmReleaseReceiveBuffer", error);
	}
}

void
sink(pmContext *pmc)
{
	int	error, i, src;
	caddr_t	addr, buf = NULL, bp = NULL;
	size_t	length;

	for (;;) {
		while ((error = pmReceive(pmc, &addr, &length)) == ENOBUFS)
			;
		if (error != PM_SUCCESS)
			fatal("pmReceive", error);
		src = ntohl(*(int *)addr);
		if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
			fatal("pmReleaseReceiveBuffer", error);
	}
}

void
burst(pmContext *pmc, int dest, size_t length, int n)
{
	int	i, error, len, m;
	caddr_t	addr, buf = NULL, bp = NULL, addr2;
	size_t	len2;

	if (length < 8)
		length = 8;

	for (i = 0; i < n; i++) {
		while ((error = pmGetSendBuffer(pmc, dest, &addr, length)) == ENOBUFS) {
			while ((error = pmReceive(pmc, &addr2, &len2)) == PM_SUCCESS) {
				if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
					fatal("pmReleaseReceiveBuffer",
					      error);
			}
			if (error != ENOBUFS)
				fatal("pmReceive", error);
		}
		if (error != PM_SUCCESS)
			fatal("pmGetSendBuffer", error);

		*(int *)addr = htonl(self);

		if ((error = pmSend(pmc)) != PM_SUCCESS)
			fatal("pmSend", error);
		while ((error = pmReceive(pmc, &addr2, &len2)) == PM_SUCCESS) {
			if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
				fatal("pmReleaseReceiveBuffer", error);
		}
		if (error != ENOBUFS)
			fatal("pmReceive", error);
	}
	while ((error = pmIsSendDone(pmc)) == EBUSY) {
		while ((error = pmReceive(pmc, &addr2, &len2)) == PM_SUCCESS) {
			if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
				fatal("pmReleaseReceiveBuffer", error);
		}
		if (error != ENOBUFS)
			fatal("pmReceive", error);
	}
	if (error != PM_SUCCESS)
		fatal("pmIsSendDone", error);
}

enum vcmd { MLOCK, MLOCK_REPLY, MUNLOCK, MUNLOCK_REPLY };

struct vreq {
	int		node;
	int		cmd;
	pmAddrHandle	hndl;
	int		length;
};

void
vreply(pmContext *pmc)
{
	int		page_size, error, dest;
	caddr_t		addr;
	pmAddrHandle	base;
	size_t		length;
	struct vreq	*q, *p;
	caddr_t		buf, b;

	if ((buf = malloc(BUFSIZE)) == NULL)
		fatal("malloc", ENOMEM);
	page_size = getpagesize();
	b = (caddr_t)(((u_long)buf + (page_size-1)) & ~(page_size-1));
	for (;;) {
		while ((error = pmReceive(pmc, &addr, &length)) == ENOBUFS)
			;
		if (error != PM_SUCCESS)
			fatal("pmReceive", error);
		p = (struct vreq *)addr;
		dest = ntohl(p->node);
		switch (ntohl(p->cmd)) {

		case MLOCK:
			if ((error = pmMLock(pmc, dest, b, ntohl(p->length), &base)) != PM_SUCCESS)
				fatal("pmMLock", error);

			while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
				;
			if (error != PM_SUCCESS)
				fatal("pmGetSendBuffer", error);
			q = (struct vreq *)addr;
			q->node = htonl(self);
			q->cmd = htonl(MLOCK_REPLY);
			q->hndl = base;
			q->length = p->length;
			if ((error = pmSend(pmc)) != PM_SUCCESS)
				fatal("pmSend", error);
			if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
				fatal("pmReleaseReceiveBuffer", error);
			break;

		case MUNLOCK:
			if ((error = pmMUnlock(pmc, dest, b, ntohl(p->length))) != PM_SUCCESS)
				fatal("pmMUnlock", error);
			while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
				;
			if (error != PM_SUCCESS)
				fatal("pmGetSendBuffer", error);
			q = (struct vreq *)addr;
			q->node = htonl(self);
			q->cmd = htonl(MUNLOCK_REPLY);
			q->hndl = p->hndl;
			q->length = p->length;
			if ((error = pmSend(pmc)) != PM_SUCCESS)
				fatal("pmSend", error);
			if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
				fatal("pmPutReceiveBuffer", error);
			break;
		}
	}
}

void
vwrite(pmContext *pmc, int dest, size_t length, int n)
{
	int		page_size, error, i, l, len, m;
	caddr_t		addr;
	size_t		len2;
	pmAddrHandle	base, hndl;
	struct vreq	*q;
	caddr_t		buf, b;

	if ((buf = malloc(BUFSIZE)) == NULL)
		fatal("malloc", ENOMEM);
	page_size = getpagesize();
	b = (caddr_t)(((u_long)buf + (page_size-1)) & ~(page_size-1));

	while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmGetSendBuffer", error);
	q = (struct vreq *)addr;
	q->node = htonl(self);
	q->cmd = htonl(MLOCK);
	q->length = htonl(length);
	if ((error = pmSend(pmc)) != PM_SUCCESS)
		fatal("pmSend", error);

	while ((error = pmReceive(pmc, &addr, &len2)) == ENOBUFS)
		;

	if (error != PM_SUCCESS)
		fatal("pmReceive", error);
	hndl = ((struct vreq *)addr)->hndl;
	if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
		fatal("pmReleaseReceiveBuffer", error);

	if ((error = pmMLock(pmc, dest, b, length, &base)) != PM_SUCCESS)
		fatal("pmMLock", error);
	for (i = 0; i < n; i++) {
		for (m = 0, len = length; len > 0; m += l, len -= l) {
			if ((l = len) > PM_RMA_MTU)
				l = PM_RMA_MTU;
			while ((error =
			    pmWrite(pmc, dest, hndl + m, base + m, l)) == ENOBUFS)
				;
			if (error != PM_SUCCESS)
				fatal("pmWrite", error);
		}
	}
	while ((error = pmIsWriteDone(pmc)) == EBUSY)
		;
	if (error != PM_SUCCESS)
		fatal("pmIsWriteDone", error);

	if ((error = pmMUnlock(pmc, dest, b, length)) != PM_SUCCESS)
		fatal("pmMUnlock", error);

	while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmGetSendBuf", error);
	q = (struct vreq *)addr;
	q->node = htonl(self);
	q->cmd = htonl(MUNLOCK);
	q->hndl = hndl;
	q->length = htonl(length);
	if ((error = pmSend(pmc)) != PM_SUCCESS)
		fatal("pmSend", error);
	while ((error = pmReceive(pmc, &addr, &length)) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmReceive", error);
	if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
		fatal("pmReleaseReceiveBuffer", error);
}

void
vread(pmContext *pmc, int dest, size_t length, int n)
{
	int		page_size, error, i, l, len, m;
	caddr_t		addr;
	size_t		len2;
	pmAddrHandle	base, hndl;
	struct vreq	*q;
	caddr_t		buf, b;

	if ((buf = malloc(BUFSIZE)) == NULL)
		fatal("malloc", ENOMEM);
	page_size = getpagesize();
	b = (caddr_t)(((u_long)buf + (page_size-1)) & ~(page_size-1));

	while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmGetSendBuffer", error);
	q = (struct vreq *)addr;
	q->node = htonl(self);
	q->cmd = htonl(MLOCK);
	q->length = htonl(length);
	if ((error = pmSend(pmc)) != PM_SUCCESS)
		fatal("pmSend", error);

	while ((error = pmReceive(pmc, &addr, &len2)) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmReceive", error);
	hndl = ((struct vreq *)addr)->hndl;
	if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
		fatal("pmReleaseReceiveBuffer", error);

	if ((error = pmMLock(pmc, dest, b, length, &base)) != PM_SUCCESS)
		fatal("pmMLock", error);
	for (i = 0; i < n; i++) {
		for (m = 0, len = length; len > 0; m += l, len -= l) {
			if ((l = len) > PM_RMA_MTU)
				l = PM_RMA_MTU;
			while ((error =
			    pmRead(pmc, dest, hndl + m, base + m, l)) == ENOBUFS)
				;
			if (error != PM_SUCCESS)
				fatal("pmRead", error);
		}
	}
	while ((error = pmIsReadDone(pmc)) == EBUSY)
		;
	if (error != PM_SUCCESS)
		fatal("pmIsReadDone", error);

	if ((error = pmMUnlock(pmc, dest, b, length)) != PM_SUCCESS)
		fatal("pmMUnlock", error);

	while ((error = pmGetSendBuffer(pmc, dest, &addr, sizeof (struct vreq))) == ENOBUFS)
		;

	if (error != PM_SUCCESS)
		fatal("pmGetSendBuffer", error);
	q = (struct vreq *)addr;
	q->node = htonl(self);
	q->cmd = htonl(MUNLOCK);
	q->hndl = hndl;
	q->length = htonl(length);
	if ((error = pmSend(pmc)) != PM_SUCCESS)
		fatal("pmSend", error);
	while ((error = pmReceive(pmc, &addr, &length)) == ENOBUFS)
		;
	if (error != PM_SUCCESS)
		fatal("pmReceive", error);
	if ((error = pmReleaseReceiveBuffer(pmc)) != PM_SUCCESS)
		fatal("pmReleaseReceiveBuffer", error);
}

size_t check_mtu(size_t length) {
	size_t mtu;

	if(pmGetMtu(pmc, dest, &mtu)) {
		mtu = 8192;
	}
	if(length > mtu) {
		fprintf(stderr, "** Specified length(%d) is larger than PM_MTU(%d) set to %d\n", length, mtu, mtu );
		fflush(stderr);
		length = mtu;
	}
	return length;
}

void
doit(pmContext *pmc, int dest, int pe, size_t *length, int iteration, int op)
{
	switch (op) {

	case PINGPONG:
		*length =  check_mtu(*length);
		pingpong(pmc, dest, *length, iteration);
		break;

	case REPLY:
		reply(pmc);
		break;

	case BURST:
		*length =  check_mtu(*length);
		burst(pmc, dest, *length, iteration);
		break;

	case SINK:
		sink(pmc);
		break;

	case VREPLY:
		vreply(pmc);
		break;

	case VWRITE:
		vwrite(pmc, dest, *length, iteration);
		break;

	case VREAD:
		vread(pmc, dest, *length, iteration);
		break;
	}
}

int
main(int argc, char **argv)
{
	int	error;
	int	ac, i, fd, mfd;
	char	*av[8];
	pmContextConfig cnf;
	enum operation op = NOP;
	enum operation op1 = NOP;

	if (argc == 1) {
		fprintf(stderr, "%s: option\n", argv[0]);
		exit(1);
	}
	scoreinit(argc, argv);

	for (argc--, argv++; argc > 0; argc--, argv++) {
		if (strcmp(argv[0], "-ping") == 0) {
			op = PINGPONG;
			op1 = REPLY;
			continue;
		}
		if (strcmp(argv[0], "-burst") == 0) {
			op = BURST;
			op1 = SINK;
			bw_flag = 1;
			continue;
		}
		if (strcmp(argv[0], "-sendrecv") == 0) {
			op = BURST;
			op1 = BURST;
			bw_flag = 1;
			continue;
		}
		if (strcmp(argv[0], "-vwrite") == 0) {
			op = VWRITE;
			op1 = VREPLY;
			bw_flag = 1;
			continue;
		}
		if (strcmp(argv[0], "-vread") == 0) {
			op = VREAD;
			op1 = VREPLY;
			bw_flag = 1;
			continue;
		}
		if (strcmp(argv[0], "-len") == 0) {
			length = strtol(argv[1], NULL, 0);
			argc--;
			argv++;
			continue;
		}
		if (strcmp(argv[0], "-dest") == 0) {
			dest = strtol(argv[1], NULL, 0);
			argc--;
			argv++;
			continue;
		}
		if (strcmp(argv[0], "-src") == 0) {
			pe = strtol(argv[1], NULL, 0);
			argc--;
			argv++;
			continue;
		}
		if (strcmp(argv[0], "-iter") == 0) {
			iteration = strtol(argv[1], NULL, 0);
			second = 0;
			argc--;
			argv++;
			continue;
		}
		fprintf(stderr, "unknown option %s\n", argv[0]);
		exit(1);
	}

	signal(SIGHUP, sig_handler);
	signal(SIGINT, sig_handler);

	if(rank == pe) {
		gettimeofday(&start_tv, NULL);
		doit(pmc, dest, pe, &length, iteration, op);
		gettimeofday(&end_tv, NULL);

		print_time(length, iteration, &start_tv, &end_tv);
		fflush(stdout);
	}
	else if(rank == dest) {
		dest = pe;
		doit(pmc, dest, pe, &length, iteration, op1);
		while(1);
	}
	else {
		/* wait until finishing jobs */
		while(1);
	}
	sc_terminate();
}