// CDXConnection - TCP/IP Connection and communication library
// Author : Jimb Esser
// Date : 1-28-1
// Version : 1.0

#include "CDXConnection.h"
#include "bsd2win.h"
#include <assert.h>
#include <stdio.h>
#include <memory.h>

typedef struct {
	unsigned char iTag;
	int iNodeNum;
	int iLen;
} CDXCON_HEADER;

CDXConnection::CDXConnection():
	m_bAcceptingConnections(false),
	m_iNumConnections(0),
	m_bConnected(false),
	m_iNodeNum(-1),
	m_bCreatedClient(false),
	m_bCreatedServer(false),
	m_iMaxConnections(MAX_CONNECTIONS),
	m_bConnectionLost(false),
	m_ulThread(0),
	m_iNumBuffers(0),
	m_iCurBuffer(0)
{
	int i;
	for (i=0; i<MAX_CONNECTIONS; i++) {
		m_Clients[i]=0;
	}
	for (i=0; i<MAX_BUFFER; i++) {
		m_Buffer[i]=NULL;
	}
}

CDXConnection::~CDXConnection()
{
	int i;
	if (m_bCreatedServer) {
		m_bAcceptingConnections=false;
		for (i=0; i<MAX_CONNECTIONS; i++) {
			if (m_Clients[i]!=0) {
				closesocket(m_Clients[i]);
			}
		}
	} else if (m_bCreatedClient) {
		closesocket(m_clientSocket);
	}
	for (i=0; i<MAX_BUFFER; i++) {
		if (m_Buffer[i]!=NULL) {
			delete m_Buffer[i];
		}
	}

	DEINIT_SOCKETS();
}


/** Host **/

HRESULT CDXConnection::CreateHost()
{
	char MyName[256];
	struct hostent *hp;

	INIT_SOCKETS();

	memset(&m_SockAddr, 0, sizeof(struct sockaddr_in));
	if (gethostname(MyName, sizeof(MyName)) == SOCKET_ERROR)
		return -1;

	hp = gethostbyname(MyName);
	if (hp == NULL)
		return -1;
	m_SockAddr.sin_family = hp->h_addrtype;

	FD_ZERO(&m_readyfds);

	m_bCreatedServer = true;

	return 0;
}

HRESULT CDXConnection::Listen(unsigned short iPort)
{
	if (!m_bCreatedServer) {
		return -1;
	}

	m_SockAddr.sin_port = htons(iPort);

	m_listenSocket = socket(AF_INET, SOCK_STREAM, 0);
	if (m_listenSocket == INVALID_SOCKET)
		return -1;

	disableTCPdelay(m_listenSocket);


	if (bind(m_listenSocket, (struct sockaddr *)&m_SockAddr, sizeof(struct sockaddr_in)) == SOCKET_ERROR)
	{
		closesocket(m_listenSocket);
		return -1;
	}

	if (listen(m_listenSocket, 3) == SOCKET_ERROR)
		return -1;	

	// Start accepting connections
	m_bAcceptingConnections=true;

	// Put this in a thread to gather connections // make it a non-blocking socket
	m_ulThread = _beginthread(ListenThread, 0, (void *) this);


	return 0;
}


void CDXConnection::stopAcceptingConnections() {
	m_bAcceptingConnections=false;
}


void ListenThread(void * dummy) {
	CDXConnection *con = (CDXConnection *)dummy;

	SOCKET s;
	while (con->m_bAcceptingConnections) {
		s = accept(con->m_listenSocket, NULL, NULL);
		if (s != INVALID_SOCKET) {
//			printf("Accepted a connection...\n");
			if (!con->m_bAcceptingConnections) {
				closesocket(s);
				return;
			}
			// Disable naggling
			con->disableTCPdelay(s);
			// Send out node number to client
			CDXCON_HEADER head;
			head.iLen=1;
			head.iNodeNum=NODE_SERVER;
			head.iTag=TAG_NODEANNOUNCE;
			int res = send(s, (char*)&head, sizeof(head), 0);
			if ((res==SOCKET_ERROR) || (res!=sizeof(head))) {
				return;
			}
			char node[2];
			node[0] = con->m_iNumConnections;
			res = send(s, node, 1, 0);
			if ((res==SOCKET_ERROR) || (res!=1)) {
				return;
			}
			// Add to list of clients
			con->m_Clients[con->m_iNumConnections++] = s;

		}
		if (con->m_iNumConnections==con->m_iMaxConnections) {
			con->m_bAcceptingConnections=false;
			return; // Stop listening
		}
	}
}


bool CDXConnection::isDataReady() {
	SOCKET max=0;
	if (!m_bCreatedServer &&
		!m_bCreatedClient) {
		return false;
	}

	if (m_bCreatedServer && (m_iNumConnections==0)) {
		return false;
	}

	if (m_bCreatedServer) { // Host
		// Clear the set and add all clients
		FD_ZERO(&m_readyfds);
		bool bEmpty=true;
		for (int i=0; i<m_iNumConnections; i++) {
			if (m_Clients[i]!=0) {
				bEmpty=false;
				FD_SET(m_Clients[i], &m_readyfds);
				if (m_Clients[i]>max) max = m_Clients[i];
			}
		}
		if (bEmpty) {
			m_iNumConnections=0;
			return false;
		}
	} else { // Client
		FD_ZERO(&m_readyfds);
		if (m_clientSocket!=0) {
			FD_SET(m_clientSocket, &m_readyfds);
				if (m_clientSocket>max) max = m_clientSocket;
		}
	}

	struct timeval timeout;
	timeout.tv_sec=0;
	timeout.tv_usec=0;

	int res;
	if ((res = select(max+1, &m_readyfds, NULL, NULL, &timeout)) == SOCKET_ERROR) {
		// Error!
		return false;
	}
	if (res > 0) {
		return true;
	} else {
		return false;
	}
}


/** Both **/
HRESULT CDXConnection::getData(int &iNodeNum, char *data, int iMaxSize, int &iNumRead, unsigned char *iTag)
{
	// Must call isDataReady to find out what socket has data
	if (!isDataReady()) {
		return -1;
	}
	SOCKET *sd;
	if (m_bCreatedServer) { // Host
		int iReady;
		for (iReady=0; (iReady<m_iNumConnections) &&
			!FD_ISSET(m_Clients[iReady], &m_readyfds); iReady++);

		if (iReady==m_iNumConnections) {
			// This shouldn't happen, there's a logic error somewhere
			OutputDebugString("Shouldn't happen!");
			assert(false);
			return -2;
		}
		sd = &m_Clients[iReady];
		iNodeNum = iReady;
	} else { // Client
		sd = &m_clientSocket;
		iNodeNum=0;
	}
	// Read in the header : 
	CDXCON_HEADER head;
	iNumRead = readFromSocket(*sd, (char*)&head, sizeof(CDXCON_HEADER));
	if (iNumRead != sizeof(CDXCON_HEADER)) {
		// Bad header
		closesocket(*sd);
		*sd=0;
		m_bConnectionLost=true;
		m_bConnected=false;
		return -5;
	}
	if (head.iLen > iMaxSize) {
		// TODO: discard bad packet or something similar
		printf("Got too large of a packet!\n");
		return -6; // Too large of a packet
	}
	if (iTag!=NULL) {
		*iTag = head.iTag;
	}
	assert(iNodeNum == head.iNodeNum);
	// Read actual data
	iNumRead = readFromSocket(*sd, data, head.iLen);
	if ((iNumRead==SOCKET_ERROR) ||
		(iNumRead<=0) ||
		(iNumRead!=head.iLen))
	{ // Connection closed or dropped or something
		closesocket(*sd);
		*sd=0;
		m_bConnectionLost=true;
		m_bConnected=false;
		return -4;
	}
	return 0;
}

HRESULT CDXConnection::sendData(int iNodeNum, const char *data, int iSize, unsigned char iTag)
{

	if (iNodeNum==NODE_ALL) {
		for (int i=0; i<m_iNumConnections; i++) {
			if (m_Clients[i]>0) {
				int res = sendData(i, data, iSize, iTag);
				if (FAILED(res)) {
					return res;
				}
			}
		}
		return 0;
	}

	CDXCON_HEADER head;
	SOCKET *sd;
	if (m_bCreatedServer) { // Host
		sd = &m_Clients[iNodeNum];
		head.iNodeNum = NODE_SERVER; // Src
	} else if (m_bCreatedClient) { // Client
		assert(iNodeNum == NODE_SERVER);
		sd = &m_clientSocket;
		head.iNodeNum = m_iNodeNum; // Src node
	} else {
		return -1; // nothing created
	}
	// Make a header
	head.iLen = iSize;
	head.iTag = iTag;
	int res;
	// Send header
	res = send(*sd, (char*)&head, sizeof(head), 0);
	if ((res==SOCKET_ERROR) || (res!=sizeof(head))) {
		return -2;
	}
	// Send data
	res = send(*sd, data, iSize, 0);
	if ((res==SOCKET_ERROR) || (res!=iSize)) {
		return -2;
	}
	return 0;
}



/** Client **/
HRESULT CDXConnection::CreateClient() {
	INIT_SOCKETS();
	m_bCreatedClient=true;
	return 0;
}

HRESULT CDXConnection::Connect(char *addr, unsigned short iPort) {
	if (!m_bCreatedClient) {
		return -1;
	}

	struct hostent *hp;

	hp = gethostbyname(addr);
	if (hp == NULL)
		return -1;

	memset(&m_SockAddr, 0, sizeof(struct sockaddr_in));
	memcpy((char *)&m_SockAddr.sin_addr, hp->h_addr, hp->h_length);

	m_SockAddr.sin_family = hp->h_addrtype;
	m_SockAddr.sin_port = htons((u_short)iPort);

	m_clientSocket = socket(hp->h_addrtype, SOCK_STREAM, 0);
	if (m_clientSocket == INVALID_SOCKET)
		return -2;

	disableTCPdelay(m_clientSocket);

	if (connect(m_clientSocket, (struct sockaddr *)&m_SockAddr, sizeof(m_SockAddr)) == SOCKET_ERROR)
	{
		closesocket(m_clientSocket);
		return -3;
	}

	int iNodeNum, iNumRead;
	unsigned char iTag=0;
	char data[255];

	while (iTag!=TAG_NODEANNOUNCE) {
		// Receieve node num announce
		while (!isDataReady()) {
			Sleep(100);
		}
		getData(iNodeNum, data, 255, iNumRead, &iTag);
	}

	m_iNodeNum = data[0];
	m_bConnected=true;
	return 0;
}



int CDXConnection::disableTCPdelay(SOCKET s)
{
	int nodelay_enabled=1;
	int optlen=sizeof(nodelay_enabled);
	return setsockopt ( s, IPPROTO_TCP, TCP_NODELAY,(LPSTR) &nodelay_enabled, optlen ); 
}


int CDXConnection::readFromSocket(SOCKET s, char *dest, int bytesToRead)
{
	int iNumRead=0;
	while (iNumRead!=bytesToRead) {
		int amt = recv(s, dest + iNumRead, bytesToRead - iNumRead, 0);
		if (amt==0) {
			// Connection closed gracefully?
			printf("Connection closed gracefully\n");
			return 0;
		} else if (amt==SOCKET_ERROR) {
			printf("Error reading from socket! (%d)\n", WSAGetLastError());
			return 0;
		} else {
			// Good read!
			iNumRead+=amt;
			assert(iNumRead<=bytesToRead);
		}
	}
	return iNumRead;
}