#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <fcntl.h>
#include <time.h>
#include <signal.h>

#include "nicStress.h"


static char *clntName = NULL;
static volatile unsigned char threadStop = FALSE;
char log_file[255];

static void usage(void)
{
	fprintf(stderr, "\n");
	fprintf(stderr, "%s \n", clntName);
	fprintf(stderr, "Usage: %s -i IPADDR | [-h] | [-v] | [-c] | [-e] | [-s SIZE] | [-d DELAY] [-l FILENAME]\n", clntName);
	fprintf(stderr, "\t-h\tprint this message\n");
	fprintf(stderr, "\t-i\tIPv4 address\n");
	fprintf(stderr, "\t-s\tsize of buffer in bytes max = 65507\n");
	fprintf(stderr, "\t-v\tverbose\n");
	fprintf(stderr, "\t-c\tcompare data\n");
	fprintf(stderr, "\t-e\tstop on error\n");
	fprintf(stderr, "\t-d\tdelay time in mseconds\n");
	fprintf(stderr, "\t-l\tlog file\n");
	fprintf(stderr, "\n");
}

static void logError(char *ipaddr, char *msg)
{
	char command[256];

	snprintf(command, 256, "echo ERROR %s: %s >> %s", ipaddr, msg, log_file);
	system(command);
}

static unsigned char isIPv4Format(const char *str)
{
	int i, j, pos[3], len;
	char buf[IP_STR_BUF];
	unsigned int addr[4];
	unsigned char valid = TRUE;

	// Length check
	len = strlen(str);
	if (len > IP_STR_LEN || len < IP_STR_LEN_MIN)
		return FALSE;
	
	strncpy(buf, str, IP_STR_BUF);

	// Look for positions of delimiters
	for (i = 0, j = 0 ; i < len ; i++) {
		
		if ( buf[i] == '.') {
			// Exceed the limit
			if (j >= 3)
				return FALSE;
			// Record & Terminate string
			pos[j] = i;
			buf[i] = 0;
			j++;
		}
	}

	// Must have 3 dots at least or at most
	if (j != 3)
		return FALSE;

	// Convert strings
	addr[0] = strtol((const char *)&buf[0], NULL, 10);
	addr[1] = strtol((const char *)&buf[pos[0] + 1], NULL, 10);
	addr[2] = strtol((const char *)&buf[pos[1] + 1], NULL, 10);
	addr[3] = strtol((const char *)&buf[pos[2] + 1], NULL, 10);

	// Check range
	for (i = 0 ; i < 4 ; i++)
    if (addr[ i ] > 255) {
		valid = FALSE;
		break;
	}

	// Not ZERO
	snprintf(buf, IP_STR_BUF, "%d.%d.%d.%d", addr[0], addr[1], addr[2], addr[3]);
	if (strcmp(buf, str))
		valid = FALSE;

	return valid;
}

static void signalHandler(int num) {

	threadStop = TRUE;
}

int main(int argc, char **argv)
{
	int ret = 0;
	int sfd, rbytes, wbytes;
	int addr_len = sizeof(struct sockaddr_in);
	struct sockaddr_in servaddr;
	struct timeval tv;
	int verbose = 0;
	int compare = 0;
	int stopOnError = 0;
	int delay = 0;
	int c, i;
	char *sbuf = NULL;
	char *rbuf = NULL;
	char *ipaddr = NULL, ipbuf[ IP_STR_BUF + 1 ];
	unsigned int size = MAX_UDP_BUFSIZE;
	char msg[256];
	int errCount = 0;


	clntName = strdup(argv[0]);

	memset(log_file, 0, sizeof(log_file));
	sprintf(log_file, "/tmp/nicUdpStressStopOnError.log");

	// Handle parameters
	while ((c = getopt(argc, argv, "i:s:d:vhcel:")) != EOF) {
		switch (c) {
			case 'i':
				strncpy(ipbuf, optarg, sizeof(ipbuf));
				if (isIPv4Format(ipbuf) == TRUE)
					ipaddr = ipbuf;
				break;
			case 'v':
				verbose = 1;
				break;
			case 's':
				size = strtol(optarg, NULL, 10);
				break;
			case 'c':
				compare = 1;
				break;
			case 'e':
				stopOnError = 1;
				break;
			case 'd':
				delay = strtol(optarg, NULL, 10);
				break;
			case 'l':
				memset(log_file, 0, sizeof(log_file));
				sprintf(log_file, "%s", optarg);
				break;
			case 'h':
			default:
				usage();
				free(clntName);
				return -1;
		}
	}

	printf("%s\n", log_file);

	// Sanity check
	if (size <= 0 || size > MAX_UDP_BUFSIZE) {
		fprintf(stderr, "%s : Invalid size(%d)\n", clntName, size);
		free(clntName);
		return -1;
	}
	if (!ipaddr) {
		fprintf(stderr, "%s : Invalid ip addr\n", clntName);
		free(clntName);
		return -1;
	}

	// Allocate buffers
	sbuf = malloc(size);
	rbuf = malloc(size);
	if (!sbuf || !rbuf) {

		ret = -1;
		fprintf(stderr, "%s : Out of memory\n", clntName);
		goto errexit;
	}

	// Handle signal
	signal(SIGTERM, signalHandler);

	// Print Start
	printf("NIC Stress Start (IP: %s)\n", ipaddr);

	sfd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (sfd < 0) {

		fprintf(stderr, "%s : Cannot open socket\n", clntName);
		goto errexit;
	}

	tv.tv_sec = 3; /* timeout value */
	tv.tv_usec = 0;
	setsockopt(sfd, SOL_SOCKET, SO_RCVTIMEO, (char*)&tv, sizeof(struct timeval));

	// Set address
	memset(&servaddr, 0, sizeof(struct sockaddr_in));
	servaddr.sin_family = AF_INET;
	servaddr.sin_port   = htons(UDP_PORT_NUM);
	ret = inet_aton(ipaddr, &servaddr.sin_addr);
	if (ret < 0) {

		fprintf(stderr, "%s : Cannot set IP address %s\n", clntName, ipaddr);

		logError(ipaddr, "Cannot set IP address");

		goto errexit;
	}

	// Initialize buffers
	for (i = 0 ; i < size ; i++)
		sbuf[i] = (char)(time( NULL ) * rand());

	// Main loop
	while (threadStop != TRUE) {

		if (compare)
			bzero(rbuf, size);

		wbytes = sendto(sfd, sbuf, size, 0, (struct sockaddr*) &servaddr, addr_len);

		if (wbytes != size)
			fprintf(stderr, "%s : Incompletely send %d of %d bytes\n", ipaddr, wbytes, size);

		if (verbose)
			printf("send %d to %s\n", wbytes, inet_ntoa(servaddr.sin_addr));

		rbytes = recvfrom(sfd, rbuf, wbytes, MSG_WAITALL, (struct sockaddr*) &servaddr, (socklen_t*) &addr_len);

		if (verbose)
			printf("recv %d from %s\n", rbytes, inet_ntoa(servaddr.sin_addr));

		if (compare && threadStop != TRUE) {
			if (rbytes != wbytes) {
				sprintf(msg, "Incompletely recv %d of %d bytes", rbytes, wbytes);
				logError(ipaddr, msg);
				errCount++;
				if (stopOnError)
					break;
			} else if (bcmp(sbuf, rbuf, size)) {
				logError(ipaddr, "data compare mismatch");
				errCount++;
				if (stopOnError)
					break;
			}

			//if (stopOnError)
			//	break;
		}

		if (delay != 0)
			usleep(delay * 1000);
	}

errexit:
	close(sfd);

	// Free buffer
	if (sbuf)
		free(sbuf);

	if (rbuf)
		free(rbuf);

	if (clntName)
		free(clntName);

	if (errCount)
	{
		snprintf(msg, 256, "echo Ethernet: FAILED `date` IP:%s >> %s", ipaddr, log_file);
		//snprintf(msg, 256, "echo ============== Ethernet FAILED ============== >> %s", log_file);
		//snprintf(msg, 256, "echo `date` >> %s", log_file);
		//snprintf(msg, 256, "echo ======================================== >> %s", log_file);
  }
	else
  {
		snprintf(msg, 256, "echo Ethernet: PASS `date` IP:%s >> %s", ipaddr, log_file);
		//snprintf(msg, 256, "echo ============== Ethernet PASS ============== >> %s", log_file);
		//snprintf(msg, 256, "echo `date` >> %s", log_file);
		//snprintf(msg, 256, "echo ======================================== >> %s", log_file);
  }
  
	system(msg);

    return ret;
}