#include "Networking.h"

DWORD WINAPI _IoCpThread(LPVOID lpParameter)
{
	LPIOCPINFO		lpIoci			= (LPIOCPINFO)lpParameter;
	DWORD			NumberOfBytes	= 0;
	LPOVERLAPPED	lpOl			= NULL;
	LPCONNECTION	lpCon			= NULL;


	while(TRUE)
	{
		if((!GetQueuedCompletionStatus(lpIoci->hIocp,
				&NumberOfBytes, (PULONG_PTR)&lpCon, &lpOl, INFINITE) && lpOl == NULL) || lpCon == NULL)
			break;

		if(lpCon != NULL && lpOl == NULL)
		{
			switch (NumberOfBytes)
			{
			case IOCPMSG_BEGINDISCONNECT:
				BeginDisconnect(lpCon);
				WaitForDisconnect(lpCon);
				WaitForSingleObject(lpCon->hAllOperationsDone, INFINITE);
				DestroyConnection(lpCon);
				break;
			}

		}
		else
		{		
			EnterCriticalSection(&lpCon->cs);
			switch(lpCon->uState)
			{
			case CS_CONNECTED:			
				if(NumberOfBytes > 0)
				{
					if(((LPBUFFER)lpOl)->isSending)
					{
						if(!JobDone(((LPBUFFER)lpOl), NumberOfBytes, ((LPBUFFER)lpOl)->cPacketSize))
						{
							if(WSASend(lpCon->sdSocket, &((LPBUFFER)lpOl)->wsa, 1, NULL, 0, lpOl, NULL) == SOCKET_ERROR)
								if(WSAGetLastError() != WSA_IO_PENDING)
								{
									RemoveOverlappedOp(lpCon);
									lpCon->lpProc(lpCon, MSG_ERROR, NumberOfBytes, (ULONG_PTR)lpOl);	// Niečo sa stalo - ohlásiť
								}
						}
						else
						{
							RemoveOverlappedOp(lpCon);
							lpCon->lpProc(lpCon, MSG_COMPLETION, NumberOfBytes, (ULONG_PTR)lpOl);
						}
					}
					else
					{
						if(!JobDone(((LPBUFFER)lpOl), NumberOfBytes, ((LPBUFFER)lpOl)->cPacketSize))
						{
							((LPBUFFER)lpOl)->dwFlags = 0;
							if(WSARecv(lpCon->sdSocket, &((LPBUFFER)lpOl)->wsa, 1, NULL, &((LPBUFFER)lpOl)->dwFlags, lpOl, NULL) == SOCKET_ERROR)
								if(WSAGetLastError() != WSA_IO_PENDING)
								{
									RemoveOverlappedOp(lpCon);
									lpCon->lpProc(lpCon, MSG_ERROR, NumberOfBytes, (ULONG_PTR)lpOl);	// Niečo zlé sa stalo - ohlásiť
								}
						}
						else
						{
							RemoveOverlappedOp(lpCon);
							lpCon->lpProc(lpCon, MSG_COMPLETION, NumberOfBytes, (ULONG_PTR)lpOl);
						}
					}
					LeaveCriticalSection(&lpCon->cs);
				}
				else
				{				
					WSASendDisconnect(lpCon->sdSocket, NULL);
					RemoveOverlappedOp(lpCon);
					LeaveCriticalSection(&lpCon->cs);
					WaitForSingleObject(lpCon->hAllOperationsDone, INFINITE);
					DestroyConnection(lpCon);
				}			
				break;

			case CS_GOINGDOWN:
				{
					if(lpCon->lpDisconnectBuffer == (LPBUFFER)lpOl)
					{
						if(NumberOfBytes != 0)
						{
							ResetBuffer(((LPBUFFER)lpOl));
							if(!JobDone(((LPBUFFER)lpOl), 0, 0))
							{						
								((LPBUFFER)lpOl)->dwFlags	= 0;
								WSARecv(lpCon->sdSocket, &((LPBUFFER)lpOl)->wsa, 1, NULL, &((LPBUFFER)lpOl)->dwFlags, lpOl, NULL);	
							}
						}
						else
						{
							lpCon->uState = CS_DOWN;
							SetEvent(lpCon->hEvent);
							DefFree(lpOl);
						}
					}
			
				}
				LeaveCriticalSection(&lpCon->cs);
				break;
			}

		
		}
	}
		
	ExitThread(0);
}


LPIOCPINFO CreateIocpServer()
{
	LPIOCPINFO		si = (LPIOCPINFO)DefAlloc(sizeof(IOCPINFO));
	SYSTEM_INFO		sysinfo;
	INT				i;
	
	si->lpData = NULL;

	GetSystemInfo(&sysinfo);

	si->hIocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE,
		NULL, (ULONG_PTR)0, sysinfo.dwNumberOfProcessors);
	if(si->hIocp == NULL)
	{
		DefFree(si);
		return NULL;
	}	
	
	si->cNumberOfThreads = 2*(sysinfo.dwNumberOfProcessors);
	
	si->aThreads = (LPHANDLE)DefAlloc((si->cNumberOfThreads * sizeof(HANDLE)));

	for(i = 0; i < si->cNumberOfThreads; i++)
	{
		si->aThreads[i] = CreateThread(NULL, 0, _IoCpThread,
		(LPVOID)si, 0, NULL);

	}	

	si->liConnections = CreateList(NULL);

	return si;
}

LPCONNECTION AddConnection(LPIOCPINFO lpIoci, SOCKET s, LPPROTOCOLPROC lpProc, ULONG cDataSize, DWORD dwRemoteAddr, USHORT uRemotePort)
{
	LPCONNECTION lpCon = (LPCONNECTION)CreateItem(lpIoci->liConnections, sizeof(CONNECTION) + cDataSize, TRUE);
	
	lpCon->uState					= CS_CONNECTED;	
	lpCon->lpIociParent				= lpIoci;
	lpCon->lpProc					= lpProc;
	lpCon->sdSocket					= s;
	lpCon->lpProtocolData			= GetProtocolDataPtr(lpCon);
	lpCon->dwRemoteAddr				= dwRemoteAddr;
	lpCon->uRemotePort				= uRemotePort;
	lpCon->lpDisconnectBuffer		= NULL;	
	lpCon->cOverlappedOperations	= 0;

	ZeroMemory(lpCon->lpProtocolData, cDataSize);
	InitializeCriticalSection(&lpCon->cs);

	lpCon->hEvent = CreateEvent(NULL, TRUE, FALSE, NULL);	
	if(lpCon->hEvent == NULL)
	{
		RemoveItem(lpIoci->liConnections, lpCon);	
		return NULL;
	}

	lpCon->hAllOperationsDone = CreateEvent(NULL, TRUE, TRUE, NULL);
	if(lpCon->hAllOperationsDone == NULL)
	{
		CloseHandle(lpCon->hAllOperationsDone);
		RemoveItem(lpIoci->liConnections, lpCon);	
		return NULL;
	}

	if(CreateIoCompletionPort((HANDLE)s, lpIoci->hIocp, (ULONG_PTR)lpCon, 0) == NULL)
	{		
		RemoveItem(lpIoci->liConnections, lpCon);		
		return NULL;
	}
	ConfigureSocket(s);	

	
	return lpCon;
}

VOID BeginDisconnect(LPCONNECTION lpCon)
{	
	EnterCriticalSection(&lpCon->cs);
	lpCon->uState = CS_GOINGDOWN;
	WSASendDisconnect(lpCon->sdSocket, NULL);
	lpCon->lpDisconnectBuffer = CreateBuffer(1024);
	if(!JobDone(lpCon->lpDisconnectBuffer, 0, 0))
	{
		lpCon->lpDisconnectBuffer->isSending	= FALSE;
		lpCon->lpDisconnectBuffer->dwFlags		= 0;
		WSARecv(lpCon->sdSocket, &lpCon->lpDisconnectBuffer->wsa, 1,
			NULL, &lpCon->lpDisconnectBuffer->dwFlags, &lpCon->lpDisconnectBuffer->ol, NULL);
	}
	LeaveCriticalSection(&lpCon->cs);
}

VOID WaitForDisconnect(LPCONNECTION lpCon)
{
	EnterCriticalSection(&lpCon->cs);	
	while(lpCon->uState != CS_DOWN)
	{
		LeaveCriticalSection(&lpCon->cs);
		WaitForSingleObject(lpCon->hEvent, INFINITE);
		EnterCriticalSection(&lpCon->cs);
	}
	LeaveCriticalSection(&lpCon->cs);

	return;
}

VOID DestroyConnection(LPCONNECTION lpCon)
{
	EnterCriticalSection(&lpCon->cs);
	closesocket(lpCon->sdSocket);
	lpCon->lpProc(lpCon, MSG_DESTROY, 0, 0);
	LeaveCriticalSection(&lpCon->cs);
	DeleteCriticalSection(&lpCon->cs);
	CloseHandle(lpCon->hEvent);
	CloseHandle(lpCon->hAllOperationsDone);
	RemoveItem(lpCon->lpIociParent->liConnections, lpCon);	
}

VOID DestroyIocpServer(LPIOCPINFO lpIoci)
{	
	INT				i;	
	LPCONNECTION	lpCon = NULL;
	LPLISTENUM		lpLe = BeginListEnum(lpIoci->liConnections);

	while(NextListItem(lpLe, (LPVOID *)&lpCon))
		BeginDisconnect(lpCon);	
	

	ResetEnum(lpLe);

	while(NextListItem(lpLe, (LPVOID *)&lpCon))
	{
		WaitForDisconnect(lpCon);
		WaitForSingleObject(lpCon->hAllOperationsDone, INFINITE);
		DestroyConnection(lpCon);
	}

	EndListEnum(lpLe);

	DestroyList(lpIoci->liConnections);

	for(i = 0; i < lpIoci->cNumberOfThreads; i ++)
		PostQueuedCompletionStatus(lpIoci->hIocp, 0, (ULONG_PTR)NULL, NULL);

	WaitForMultipleObjects(lpIoci->cNumberOfThreads, lpIoci->aThreads, TRUE, INFINITE);

	for(i = 0; i < lpIoci->cNumberOfThreads; i ++)
		CloseHandle(lpIoci->aThreads[i]);

	CloseHandle(lpIoci->hIocp);
	

	HeapFree(g_hProcHeap, 0, lpIoci->aThreads);
	HeapFree(g_hProcHeap, 0, lpIoci);
}


VOID SendBuffer(LPCONNECTION lpCon, LPBUFFER lpBuf, ULONG cPacketSize)
{	
	lpBuf->cPacketSize	= cPacketSize;
	lpBuf->isSending	= TRUE;
	if(!JobDone(lpBuf, 0, lpBuf->cPacketSize))
		if(WSASend(lpCon->sdSocket, &lpBuf->wsa, 1, NULL, 0, &lpBuf->ol, NULL) == SOCKET_ERROR)
			if(WSAGetLastError() != WSA_IO_PENDING)
			{
				lpCon->lpProc(lpCon, MSG_ERROR, 0, (ULONG_PTR)&lpBuf->ol);
			}
			else
				AddOverlappedOp(lpCon);
}

VOID RecvBuffer(LPCONNECTION lpCon, LPBUFFER lpBuf)
{	
	lpBuf->cPacketSize	= 0;
	lpBuf->isSending	= FALSE;
	if(!JobDone(lpBuf, 0, lpBuf->cPacketSize))
	{
		lpBuf->dwFlags = 0;
		if(WSARecv(lpCon->sdSocket, &lpBuf->wsa, 1, NULL, &lpBuf->dwFlags, &lpBuf->ol, NULL) == SOCKET_ERROR)
			if(WSAGetLastError() != WSA_IO_PENDING)
			{				
				lpCon->lpProc(lpCon, MSG_ERROR, 0, (ULONG_PTR)&lpBuf->ol);
			}
			else
				AddOverlappedOp(lpCon);
	}
}

BOOL ConfigureSocket(SOCKET s)
{
	LINGER	lin;
	DWORD	dwVal;

	dwVal = 0;
	if(setsockopt(s, SOL_SOCKET, SO_RCVBUF, (char *)&dwVal, sizeof(DWORD)) == SOCKET_ERROR)
		return FALSE;

	dwVal = 0;
	if(setsockopt(s, SOL_SOCKET, SO_SNDBUF, (char *)&dwVal, sizeof(DWORD)) == SOCKET_ERROR)
		return FALSE;

	lin.l_onoff		= 0;
	lin.l_linger	= 0;
	if(setsockopt(s, SOL_SOCKET, SO_LINGER, (char *)&lin, sizeof(LINGER)) == SOCKET_ERROR)
		return FALSE;

	return TRUE;
}

SOCKET ConnectTo(LPENDPOINT lpEp)
{
	SOCKET s = INVALID_SOCKET;
	SOCKADDR_IN sin;

	ZeroMemory(&sin, sizeof(SOCKADDR_IN));
	s = WSASocket(AF_INET, SOCK_STREAM, IPPROTO_TCP, NULL, 0, WSA_FLAG_OVERLAPPED);
	if(s == INVALID_SOCKET)
		return s;

	sin.sin_family	= AF_INET;	
	WSAHtonl(s, lpEp->dwAddr, &sin.sin_addr.S_un.S_addr);
	WSAHtons(s, lpEp->uPort, &sin.sin_port);

	if(WSAConnect(s, (struct sockaddr*)&sin, sizeof(SOCKADDR_IN), NULL, NULL, NULL, NULL) == SOCKET_ERROR)
	{
		closesocket(s);
		return INVALID_SOCKET;
	}
	return s;
}

VOID AddOverlappedOp(LPCONNECTION lpCon)
{
	lpCon->cOverlappedOperations ++;
	if(lpCon->cOverlappedOperations == 1)
		ResetEvent(lpCon->hAllOperationsDone);
}

VOID RemoveOverlappedOp(LPCONNECTION lpCon)
{
	lpCon->cOverlappedOperations --;
	if(lpCon->cOverlappedOperations == 0)
		SetEvent(lpCon->hAllOperationsDone);
}