/* 
 * Copyright 2011 Stefan Lankes, Chair for Operating Systems,
 *                               RWTH Aachen University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * This file is part of MetalSVM. 
 */

#include <metalsvm/stddef.h>
#include <metalsvm/stdio.h>
#include <metalsvm/time.h>
#include <metalsvm/tasks.h>
#include <metalsvm/syscall.h>
#include <metalsvm/errno.h>

#include "tests.h"

#ifdef CONFIG_ROCKCREEK
#include <asm/RCCE.h>
#include <asm/RCCE_lib.h>
#include <asm/SCC_API.h>
#endif

#define USE_SOCKET_BYPASSING 0

/*
 * This implements a netio server and client (only TCP version).
 * The client sends a command word (4 bytes) then a data length word (4 bytes).
 * If the command is "receive", the server is to consume "data length" bytes into
 * a circular buffer until the first byte is non-zero, then it is to consume
 * another command/data pair.
 * If the command is "send", the server is to send "data length" bytes from a circular
 * buffer with the first byte being zero, until "some time" (6 seconds in the
 * current netio131.zip download) has passed and then send one final buffer with
 * the first byte being non-zero. Then it is to consume another command/data pair.
 */

/* See http://www.nwlab.net/art/netio/netio.html to get the netio tool */

#if (defined(START_NETIO) || defined(START_CHIEFTEST)) && defined(CONFIG_LWIP)
#ifdef CONFIG_ROCKCREEK
#if USE_SOCKET_BYPASSING // for socket bypassing
#include <lwip/opt.h>
#undef LWIP_COMPAT_SOCKETS
#endif
#endif

#include <lwip/sockets.h>
#include <lwip/err.h>
#include <lwip/stats.h>

#ifdef CONFIG_ROCKCREEK
#if USE_SOCKET_BYPASSING // for socket bypassing
#include <net/mmnif.h>
#undef AF_INET
#define AF_INET	AF_MMNIF_NET
#endif
#endif

typedef struct
{
	uint32_t cmd;
	uint32_t data;
} CONTROL;

#define CMD_QUIT  0
#define CMD_C2S   1
#define CMD_S2C   2
#define CMD_RES   3

#define CTLSIZE sizeof(CONTROL)
#define DEFAULTPORT 0x494F
#define TMAXSIZE 65536

static int tSizes[] = {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32767};
static size_t ntSizes = sizeof(tSizes) / sizeof(int);
static int nPort = DEFAULTPORT;
static const int sobufsize = 131072;
static struct in_addr addr_local;
static struct in_addr addr_server;

static int send_data(int socket, void *buffer, size_t size, int flags)
{
	ssize_t rc = send(socket, buffer, size, flags);

	if (rc < 0)
	{
		kprintf("send failed: %d\n", errno);
		return -1;
	}

	if (rc != size)
		return 1;

	return 0;
}

static int recv_data(int socket, void *buffer, size_t size, int flags)
{
	ssize_t rc = recv(socket, buffer, size, flags);

	if (rc < 0) {
		kprintf("recv failed: %d\n", errno);
		return -1;
	}

	if (rc != size)
		return 1;

	return 0;
}

static char *InitBuffer(size_t nSize)
{
	char *cBuffer = kmalloc(nSize); 

	memset(cBuffer, 0xFF, nSize); 
	cBuffer[0] = 0;

	return cBuffer;
}

static char *PacketSize(int nSize)
{
	static char szBuffer[64];

	if ((nSize % 1024) == 0 || (nSize % 1024) == 1023)
		ksprintf(szBuffer, "%2dk", (nSize + 512) / 1024);
	else
		ksprintf(szBuffer, "%d", nSize);

	return szBuffer;
}

static int TCPServer(void* arg)
{
	char *cBuffer;
	CONTROL ctl;
	uint64_t nData;
	struct sockaddr_in sa_server, sa_client;
	int server, client;	
	socklen_t length;
	struct timeval tv;
	fd_set fds;
	int rc;
	int nByte;
	int err;
	uint64_t start, end;
	uint32_t freq = get_cpu_frequency(); /* in MHz */

	if ((cBuffer = InitBuffer(TMAXSIZE)) == NULL) {
    		kprintf("Netio: Not enough memory\n");
		return -1;
	}

	if ((server = socket(PF_INET, SOCK_STREAM, 0)) < 0) {
		kprintf("socket failed: %d\n", server);
 		kfree(cBuffer, TMAXSIZE);
 		return -1;
	}

	setsockopt(server, SOL_SOCKET, SO_RCVBUF, (char *) &sobufsize, sizeof(sobufsize));
	setsockopt(server, SOL_SOCKET, SO_SNDBUF, (char *) &sobufsize, sizeof(sobufsize));

	memset((char *) &sa_server, 0x00, sizeof(sa_server));
	sa_server.sin_family = AF_INET;
	sa_server.sin_port = htons(nPort);
	sa_server.sin_addr = addr_local;

	if ((err = bind(server, (struct sockaddr *) &sa_server, sizeof(sa_server))) < 0)
	{
		kprintf("bind failed: %d\n", errno);
		closesocket(server);
		kfree(cBuffer, TMAXSIZE);
		return -1;
	}

	if ((err = listen(server, 2)) != 0)
	{
		kprintf("listen failed: %d\n", errno);
		closesocket(server);
		kfree(cBuffer, TMAXSIZE);
		return -1;
	}

	for (;;)
	{
		kprintf("TCP server listening.\n");

#if !USE_SOCKET_BYPASSING
		FD_ZERO(&fds);
		FD_SET(server, &fds);
		tv.tv_sec  = 3600;
		tv.tv_usec = 0;

		if ((rc = select(FD_SETSIZE, &fds, 0, 0, &tv)) < 0)
		{
			kprintf("select failed: %d\n", errno);
			break;
		}

		if (rc == 0 || FD_ISSET(server, &fds) == 0)
			continue;
#endif
		length = sizeof(sa_client);
#if USE_SOCKET_BYPASSING
		// TODO: Bug, not compatible with BSD sockets
		memcpy(&sa_client, &sa_server, length);
#endif
		if ((client = accept(server, (struct sockaddr *) &sa_client, &length)) < 0) {
			kprintf("accept faild: %d\n", errno);
			continue;
		}

		setsockopt(client, SOL_SOCKET, SO_RCVBUF, (char *) &sobufsize, sizeof(sobufsize));
		setsockopt(client, SOL_SOCKET, SO_SNDBUF, (char *) &sobufsize, sizeof(sobufsize));

		kprintf("TCP connection established ... ");

		for (;;)
		{
			if (recv_data(client, (void *) &ctl, CTLSIZE, 0))
				break;

			ctl.cmd = ntohl(ctl.cmd);
			ctl.data = ntohl(ctl.data);

			if (ctl.cmd == CMD_C2S)
			{
				start = rdtsc();

				kprintf("\nReceiving from client, packet size %s ... \n", PacketSize(ctl.data));
				cBuffer[0] = 0;
				nData = 0;
	
				do {
					for (nByte = 0; nByte < ctl.data; )
					{
						rc = recv(client, cBuffer + nByte, ctl.data - nByte, 0);

						if (rc < 0)
						{
							kprintf("recv failed: %d\n", errno);
							break;
						}
            
						if (rc > 0)
							nByte += rc;
					}

					nData += ctl.data;
				} while (cBuffer[0] == 0 && rc > 0);

				end = rdtsc();
				kprintf("Time to receive %llu bytes: %llu nsec (ticks %llu)\n", nData, ((end-start)*1000ULL)/freq, end-start);
			} else if (ctl.cmd == CMD_S2C) {
				start = rdtsc();

				kprintf("\nSending to client, packet size %s ... \n", PacketSize(ctl.data));
				cBuffer[0] = 0;
				nData = 0;

				do
				{
					//GenerateRandomData(cBuffer, ctl.data);

					for (nByte = 0; nByte < ctl.data; )
					{
						rc = send(client, cBuffer + nByte, ctl.data - nByte, 0);

						if (rc < 0)
						{
							kprintf("send failed: %d\n", errno);
							break;
						}

						if (rc > 0)
							nByte += rc;
					}

					nData += ctl.data;
					end = rdtsc();
				} while((end-start)/freq < 6000000ULL /* = 6s */);

				cBuffer[0] = 1;

				if (send_data(client, cBuffer, ctl.data, 0))
					break;

				end = rdtsc();
				kprintf("Time to send %llu bytes: %llu nsec (ticks %llu)\n", nData, ((end-start)*1000ULL)/freq, end-start);
			} else /* quit */
				break;
		}

		kprintf("\nDone.\n");

		closesocket(client);

		if (rc < 0)
			break;
	}

	closesocket(server);
	kfree(cBuffer, TMAXSIZE);

	return 0;
}

int TCP_Bench(void)
{
	char *cBuffer;
	CONTROL ctl;
	uint64_t nData;
	int i;
	struct sockaddr_in sa_server;
	int server;
	int rc, err;
	int nByte;
	uint64_t start, end;
	uint32_t freq = get_cpu_frequency(); /* in MHz */

	if ((cBuffer = InitBuffer(TMAXSIZE)) == NULL)
	{
		kprintf("Netio: Not enough memory\n");
		return -1;
	}

	if ((server = socket(PF_INET, SOCK_STREAM, 0)) < 0)
	{
		kprintf("socket failed: %d\n", errno);
		kfree(cBuffer, TMAXSIZE);
		return -2;
	}

	setsockopt(server, SOL_SOCKET, SO_RCVBUF, (char *) &sobufsize, sizeof(sobufsize));
	setsockopt(server, SOL_SOCKET, SO_SNDBUF, (char *) &sobufsize, sizeof(sobufsize));

	sa_server.sin_family = AF_INET;
	sa_server.sin_port = htons(nPort);
	sa_server.sin_addr = addr_server;

	if ((err = connect(server, (struct sockaddr *) &sa_server, sizeof(sa_server))) < 0)
	{
		kprintf("connect failed: %d\n", errno);
		closesocket(server);
		kfree(cBuffer, TMAXSIZE);
		return -2;
	}

	kprintf("\nTCP connection established.\n");

	for (i = 0; i < ntSizes; i++)
	{
		kprintf("Packet size %s bytes: ", PacketSize(tSizes[i]));

		/* tell the server we will send it data now */

		ctl.cmd = htonl(CMD_C2S);
		ctl.data = htonl(tSizes[i]);

		if (send_data(server, (void *) &ctl, CTLSIZE, 0))
			break;

		/* 1 - Tx test */

		start = rdtsc();
		nData = 0;
		cBuffer[0] = 0;

		do
		{
			//GenerateRandomData(cBuffer, tSizes[i]);

			for (nByte = 0; nByte < tSizes[i]; )
			{
				rc = send(server, cBuffer + nByte, tSizes[i] - nByte, 0);

				if (rc < 0)
				{
					kprintf("send failed: %d\n", errno);
					return -1;
				}
          
				if (rc > 0)
					nByte += rc;
			}

			nData += tSizes[i];
			end = rdtsc();
		} while((end-start)/freq < 6000000ULL /* = 6s */);

		kprintf("%llu/100 MBytes/s", ((100ULL*nData)/(1024ULL*1024ULL))/((end-start)/(1000000ULL*freq)));

		kprintf(" Tx, ");

		cBuffer[0] = 1;

		if (send_data(server, cBuffer, tSizes[i], 0))
			break;

		/* tell the server we expect him to send us data now */

		ctl.cmd = htonl(CMD_S2C);
		ctl.data = htonl(tSizes[i]);

		if (send_data(server, (void *) &ctl, CTLSIZE, 0))
			break;

		/* 2 - Rx test */

		start = rdtsc();
		nData = 0;
		cBuffer[0] = 0;
		rc = 0;

		do
		{
			for (nByte = 0; nByte < tSizes[i]; )
			{
				rc = recv(server, cBuffer + nByte, tSizes[i] - nByte, 0);

				if (rc < 0)
				{
					kprintf("recv failed: %d\n", errno);
					return -1;
				}
          
				if (rc > 0)
					nByte += rc;
			}

			nData += tSizes[i];
		} while (cBuffer[0] == 0 && rc > 0);

		end = rdtsc();
		kprintf("%llu/100 MBytes/s", ((100ULL*nData)/(1024ULL*1024ULL))/((end-start)/(1000000ULL*freq)));

		kprintf(" Rx.\n");
	}

	ctl.cmd = htonl(CMD_QUIT);
	ctl.data = 0;

	send_data(server, (void *) &ctl, CTLSIZE, 0);

	kprintf("Done.\n");

	closesocket(server);
	kfree(cBuffer, TMAXSIZE);

	return 0;
}

int netio_init(void)
{
	int err = 0;

	addr_local.s_addr = INADDR_ANY;
	//addr_server.s_addr = inet_addr("192.168.28.254"); 
	addr_server.s_addr = inet_addr("192.168.0.3");

#ifdef CONFIG_ROCKCREEK
	if (RCCE_ue() == 2) {
		err = create_kernel_task(NULL, TCPServer, NULL, NORMAL_PRIO);
	} else if (RCCE_ue() == 0) {
		sleep(3);
		err = TCP_Bench();
		stats_display();
	}
#else
	err = create_kernel_task(NULL, TCPServer, NULL, NORMAL_PRIO);
#endif

	return err;
}
#endif