/*
 * thrulay_tu.c -- a thrulay library test utility.
 * 
 * Written by Huadong Liu, http://www.cs.utk.edu/~hliu/
 * 
 * Copyright 2005, Internet2.
 * Legal conditions are in file LICENSE
 * (MD5 = ecfa50d1b0bfbb81b658c810d0476a52).
 */

#include <stdio.h>
#include <stdlib.h>
#include <limits.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include "thrulay.h"

#define MAXLINE 128

/*
 * Type of tests.
 */ 
typedef enum _test_type {
	TEST_INIT = 0,
	TEST_TCP_PROG_NORMAL, /* Progressively report test status. */
	TEST_TCP_PROG_STOP,   /* Progressively report test status, stop in middle */
	TEST_TCP_LAZY_NORMAL, /* Report test status when test exits. -- TCP */
	TEST_TCP_LAZY_STOP,   /* Wait half of test duration then cancel. -- TCP */
	TEST_UDP_LAZY_NORMAL, /* Report test status when test exits. -- UDP */
	TEST_UDP_LAZY_STOP,   /* Wait half of test duration then cancel. -- UDP*/
	TEST_EXIT
} test_type_t;

/*
 * Print test menu.
 */ 
void 
print_menu()
{
	fprintf(stderr, "\n");
	fprintf(stderr, "1. TCP PROGRESSIVE TEST -- NORMAL\n");
	fprintf(stderr, "2. TCP PROGRESSIVE TEST -- STOP IN MIDDLE\n");
	fprintf(stderr, "3. TCP LAZY TEST -- NORMAL\n");
	fprintf(stderr, "4. TCP LAZY TEST -- STOP IN MIDDLE\n");
	fprintf(stderr, "5. UDP LAZY TEST --NORMAL\n");
	fprintf(stderr, "6. UDP LAZY TEST -- STOP IN MIDDLE\n");
	fprintf(stderr, "7. EXIT\n");
	fprintf(stderr, "Your choice:");
}

/* 
 * Convert a string like "12k" to a number like 12000. Return zero on error.
 */
u_int64_t
rate2i(char *s)
{
	u_int64_t r;    /* Result. */
	char *p;
	int l;
	int suffix = 0;

	/* First, set the multiple. */
	l = strlen(s);
	switch (s[l-1]) {
	case '0':
	case '1':
	case '2':
	case '3':
	case '4':
	case '5':
	case '6':
	case '7':
	case '8':
	case '9':
		r = 1;
		break;
	case 'k':
	case 'K':
		r = 1000;
		suffix = 1;
		break;
	case 'm':
	case 'M':
		r = 1000000;
		suffix = 1;
		break;
	case 'g':
	case 'G':
		r = 1000000000ULL;
		suffix = 1;
		break;
	case 't':
	case 'T':
		r = 1000000000000ULL;
		suffix = 1;
		break;
	default:
		r = 0;
	}
	if (suffix)
		s[l-1] = '\0';
	r *= strtoll(s, &p, 10);
	if (!*s || *p)  /* Invalid character found. */
		r = 0;
	return r;
}

/*
 * Get an integer in range (min, max) from stdin. Allow default value by
 * simply hit Enter.
 */ 
int 
get_int(int min, int max)
{
	char buffer[MAXLINE] = {0};
	int number;
	char *p;

	while (1) {
		fgets(buffer, MAXLINE, stdin); 
		p = strchr(buffer, '\n');
		if (p == NULL) {
			fprintf(stderr, "Your choice is too long!\n");
			exit(0);
		}
		if (p == buffer) return 0;

		*p = '\0';
		number = atoi(buffer);
		if ((number > max) || (number < min)) {
			fprintf(stderr, "Invalid number, do it one more time\n");
			continue;
		}
		return number;
	}
}

/*
 * Get a string from stdin. If empty is allowed, returns when Enter is hit.
 */ 
void
get_str(char *buffer, int allow_empty)
{
	char *p;
	
	while (1) {
		fgets(buffer, MAXLINE, stdin); 
		p = strchr(buffer, '\n');
		if (p == NULL) {
			fprintf(stderr, "Your choice is too long!\n");
			exit(0);
		}
		if ((!allow_empty) && (p == buffer)) {
			fprintf(stderr, "Empty string, do it one more time\n");
		}
		else {
			*p = '\0';
			return;
		}
	}
}

/*
 * Read setting of a thrulay TCP test.
 */ 
void 
read_tcp_setting(thrulay_setting_tcp_t *setting_tcp)
{
	int rc;
	char host[MAXLINE] = {0};
	
	fprintf(stderr, "Server name:");
	get_str(host, 0);
	setting_tcp->server_name = strdup(host);

	fprintf(stderr, "Server port [5003]:");
	rc = get_int(5003, 9999);
	setting_tcp->server_port = rc ? rc : 5003;
		
	fprintf(stderr, "Test duration [60s]:");
	rc = get_int(1, 24 * 60 * 60);
	setting_tcp->test_duration = rc ? rc : 60;
		
	fprintf(stderr, "Report_interval [1s]:");
	rc = get_int(1, setting_tcp->test_duration);
	setting_tcp->report_interval = rc ? rc : 1;
	
	fprintf(stderr, "Window size [4194304B]:");
	rc = get_int(1500, INT_MAX);
	setting_tcp->window_size = rc ? rc : 4194304;
	
	fprintf(stderr, "Block size [8192B]:");
	rc = get_int(16, 1048576);
	setting_tcp->block_size = rc ? rc : 8192;

	fprintf(stderr, "Number of parallel streams [1]:");
	rc = get_int(1, 64);
	setting_tcp->nstream = rc ? rc : 1;

	/* We do not check TOS byte here. If you want to do so, see thrulay.c. */
	fprintf(stderr, "TOS byte [0]:");
	rc = get_int(1, 255);
	setting_tcp->tos = rc ? rc : 0;
	
	fprintf(stderr, "\n");
}

/*
 * Read setting of a thrulay UDP test.
 */ 
void 
read_udp_setting(thrulay_setting_udp_t *setting_udp)
{
	int rc;
	char buffer[MAXLINE] = {0};
	
	fprintf(stderr, "Server name:");
	get_str(buffer, 0);
	setting_udp->server_name = strdup(buffer);

	fprintf(stderr, "Server port [5003]:");
	rc = get_int(5003, 9999);
	setting_udp->server_port = rc ? rc : 5003;
		
	fprintf(stderr, "Test duration [60s]:");
	rc = get_int(1, 24 * 60 * 60);
	setting_udp->test_duration = rc ? rc : 60;
		
	fprintf(stderr, "UDP buffer size [4194304B]:");
	rc = get_int(1500, INT_MAX);
	setting_udp->udp_buffer_size = rc ? rc : 4194304;
	
	fprintf(stderr, "Packet size [1500B]:");
	rc = get_int(20+8, 1048576);
	setting_udp->packet_size = rc ? rc : 1500;

	fprintf(stderr, "Number of parallel streams [1]:");
	rc = get_int(1, 64);
	setting_udp->nstream = rc ? rc : 1;
	
	fprintf(stderr, "TOS byte [0]:");
	rc = get_int(1, 255);
	setting_udp->tos = rc ? rc : 0;

	setting_udp->rate = 0;
	while (!setting_udp->rate) {
		fprintf(stderr, "Sending rate (KMGT/s):");
		get_str(buffer, 0);
		setting_udp->rate = rate2i(buffer); 
	}
		
	fprintf(stderr, "Busy wait (yes/no) [yes]:");
	get_str(buffer, 1);
	if ((buffer[0] = 'n') || (buffer[0] = 'N')) {
		setting_udp->wt = THRULAY_NONBUSYWAIT;
	}
	else {
		setting_udp->wt = THRULAY_BUSYWAIT;
	}

	fprintf(stderr, "\n");
}

/*
 * Show how to get a thrulay report. You can further print individual 
 * report fields as in thrulay.c.
 */ 
void 
print_report(int tid, int nstream, int from, int to)
{
	int i, j;
	thrulay_report_tcp_t report;

	for (i=0; i<nstream; i++) {
		for (j=from; j<to; j++) {
			if (0 != thrulay_get_report(tid, i, j+1, &report)) {
				fprintf(stderr, "cannot get the %dth report "
						"of the %d-th stream\n", j+1, i+1);
			}
			else {
				fprintf(stderr, "get report[%d] of stream[%d]\n", 
						j+1, i+1);
			} /* if */
		} /* for j */
	} /* for i */

	return;
}

/*
 * Do a progressive TCP test. If half_stop is set, the test will be
 * terminated in the middle. Note that you do not have to close the
 * test before you call thrulay_exit(). thrulay_exit() will do this
 * for you.
 */ 
void
perform_prog_test(thrulay_setting_tcp_t *setting_tcp, int half_stop)
{
	int rc;
	int tid;
	int status;
	int already_got = 0;
	int elapsed = 0;

	rc = thrulay_init(); /* Initialize the thrulay library. */
	if (rc < 0) {
		thrulay_err_msg(rc);
		fprintf(stderr, "could not initialize thrulay lib\n");
		return;
	}
	else {
		fprintf(stderr, "initialize thrulay library successfully\n");
	}

	tid = thrulay_open(THRULAY_TCP, setting_tcp); /* Setup the test. */
	if (tid < 1) {
		thrulay_err_msg(tid);
		fprintf(stderr, "open a TCP test failed\n");
		thrulay_exit();
		return;
	}
	else {
		fprintf(stderr, "open a test successfully, tid = %d\n", tid);
	}

	rc = thrulay_start(tid); /* Start the test. */
	if (tid != rc) {
		thrulay_err_msg(rc);
		fprintf(stderr, "start a test %d failed\n", tid);
		if (0 != (rc=thrulay_close(tid))) {
			thrulay_err_msg(rc);
			fprintf(stderr, "cannot close a test %d\n", tid);
		}
		thrulay_exit();
		return;
	}

	sleep(setting_tcp->report_interval);
	do {
		rc = thrulay_wait(tid, &status, THRULAY_POLL);
		if (rc < 0) {
			thrulay_err_msg(rc);
			fprintf(stderr, "poll test status failed\n");
			thrulay_exit();
			return;
		}
		else if (rc == 0) {
			if (status == 0) {
				usleep(1000);
				continue;
			}
			else {
				print_report(tid, setting_tcp->nstream, already_got, status);
				already_got = status;
			}
			elapsed += setting_tcp->report_interval;
			if (half_stop && (elapsed >= setting_tcp->test_duration/2)) break;
			sleep(setting_tcp->report_interval);
		}
		else {
			if (status > 0) {
				print_report(tid, setting_tcp->nstream, already_got, status);
			}
			else if (status < 0) {
				thrulay_err_msg(status);
				fprintf(stderr, "test %d exit with error\n", tid);
			}
		}
	} while (rc != tid);
	
	if (half_stop) {
		rc = thrulay_wait(tid, &status, THRULAY_STOP);
		if (tid != rc) {
			thrulay_err_msg(rc);
			fprintf(stderr, "stop test %d failed\n", tid);
			if (0 != thrulay_close(tid)) {
				fprintf(stderr, "cannot close a test %d\n", tid);
			}
			thrulay_exit();
			return;
		}
		else {
			fprintf(stderr, "test %d stopped\n", tid);
			print_report(tid, setting_tcp->nstream, already_got, status);
		}
	}
	
	thrulay_exit(); /* exit() should call close() */
	fprintf(stderr, "cleanup thrulay libray successfully\n");
}

void
perform_lazy_test(thrulay_traffic_type_t type, void *setting, int half_stop)
{
	int tid;                            /* unique test id */    
	int rc;
	int status;
	int nstream;
	int i, j;
	void *report;
	thrulay_report_tcp_t rpt_tcp;
	thrulay_report_udp_t rpt_udp;
	thrulay_setting_tcp_t *setting_tcp;
	thrulay_setting_udp_t *setting_udp;

	if (type == THRULAY_TCP) {
		setting_tcp = (thrulay_setting_tcp_t *)setting;
		nstream = setting_tcp->nstream;
	}
	else {
		setting_udp = (thrulay_setting_udp_t *)setting;
		nstream = setting_udp->nstream;
	}
	
	rc = thrulay_init(); /* Initialize the thrulay library. */
	if (rc < 0) {
		thrulay_err_msg(rc);
		fprintf(stderr, "could not initialize thrulay lib\n");
		return;
	}
	else {
		fprintf(stderr, "initialize thrulay library successfully\n");
	}

	tid = thrulay_open(type, setting); /* Setup the test. */
	if (tid < 1) {
		thrulay_err_msg(tid);
		fprintf(stderr, "open a test failed\n");
		thrulay_exit();
		return;
	}
	else {
		fprintf(stderr, "open a test successfully, tid = %d\n", tid);
	}
	
	rc = thrulay_start(tid); /* Start the test. */
	if (tid != rc) {
		thrulay_err_msg(rc);
		fprintf(stderr, "start a test %d failed\n", tid);
		if (0 != thrulay_close(tid)) {
			fprintf(stderr, "cannot close a test %d\n", tid);
		}
		thrulay_exit();
		return;
	}

	if (half_stop) {	
		if (type == THRULAY_TCP) {
			sleep(setting_tcp->test_duration/2);
		}
		else {
			sleep(setting_udp->test_duration/2);
		}

		rc = thrulay_wait(tid, &status, THRULAY_STOP);
		if (tid != rc) {
			fprintf(stderr, "stop test %d failed\n", tid);
			if ((rc=thrulay_close(tid)) != 0) {
				thrulay_err_msg(rc);
				fprintf(stderr, "cannot close a test %d\n", tid);
			}
			thrulay_exit();
			return;
		}
		else {
			fprintf(stderr, "test %d stopped, # of reports=%d\n", tid, status);
		}
	}
	else {
		rc = thrulay_wait(tid, &status, THRULAY_WAIT);
		if (tid != rc) {
			thrulay_err_msg(rc);
			fprintf(stderr, "wait for test %d failed\n", tid);
			if ((rc=thrulay_close(tid)) != 0) {
				thrulay_err_msg(rc);
				fprintf(stderr, "cannot close a test %d\n", tid);
			}
			thrulay_exit();
			return;
		}
		else {
			fprintf(stderr, "test %d finished, # of reports=%d\n", tid, status);
		}
	}

	for (i=0; i<nstream; i++) {
		for (j=1; j<=status; j++) {
			if (type == THRULAY_TCP) {
				report = &rpt_tcp;
			}
			else {
				report = &rpt_udp;
			}
			if (0 != (rc=thrulay_get_report(tid, i, j, report))) {
				thrulay_err_msg(rc);
				fprintf(stderr, "cannot get the %dth report of the %d-th" 
						"stream\n", j, i+1);
			}
			else {
				fprintf(stderr, "get report[%d] of stream[%d]\n", j, i+1);
			}
		}
	}

	thrulay_exit(); /* exit() should call close() */
	fprintf(stderr, "cleanup thrulay libray successfully\n");
}

int 
main(int argc, char *argv[])
{
	thrulay_setting_tcp_t setting_tcp;   /* settings for TCP test */
	thrulay_setting_udp_t setting_udp;   /* settings for UDP test */
	test_type_t nextstep = TEST_INIT;

	while (nextstep != TEST_EXIT) {
		switch (nextstep) {
		case TEST_INIT:
			print_menu();
			nextstep = get_int(TEST_TCP_PROG_NORMAL, TEST_EXIT);
			break;
		case TEST_TCP_PROG_NORMAL:
			read_tcp_setting(&setting_tcp);
			perform_prog_test(&setting_tcp, 0);
			nextstep = TEST_INIT;
			break;
		case TEST_TCP_PROG_STOP:
			read_tcp_setting(&setting_tcp);
			perform_prog_test(&setting_tcp, 1);
			nextstep = TEST_INIT;
			break;
		case TEST_TCP_LAZY_NORMAL:
			read_tcp_setting(&setting_tcp);
			perform_lazy_test(THRULAY_TCP, &setting_tcp, 0);
			nextstep = TEST_INIT;
			break;
		case TEST_TCP_LAZY_STOP:
			read_tcp_setting(&setting_tcp);
			perform_lazy_test(THRULAY_TCP, &setting_tcp, 1);
			nextstep = TEST_INIT;
			break;
		case TEST_UDP_LAZY_NORMAL:
			read_udp_setting(&setting_udp);
			perform_lazy_test(THRULAY_UDP, &setting_udp, 0);
			nextstep = TEST_INIT;
			break;
		case TEST_UDP_LAZY_STOP:
			read_udp_setting(&setting_udp);
			perform_lazy_test(THRULAY_UDP, &setting_udp, 1);
			nextstep = TEST_INIT;
			break;
		default:
			break;
		}
	}

	exit(0);
}