//
// Asura_TCPSocket.cpp
//
// Created by KOR - 10 February 2009
//
// Copyright (C) 2009 Rebellion, All Rights Reserved.
//

#include "Core/Asura.h"

#ifdef ASURA_SUPPORT_EXTERNAL_NETWORK

// This file's header.
#include "Network/Asura_TCPSocket.h"

#ifdef _DEBUG
// Includes for unit test
#include "Core/Asura_Timers.h"
#include "Core/Asura_Multithread.h"
#endif
#include "Core/Asura_Multithread.h"
#include "Debugging/Asura_Logfile.h"
#include "Network/Asura_ReliableSocket.h"

/// Macros ////////////////////////////////////////////////////////////////////////////////////////
#ifdef ASURA_TCP_LOGGING_ENABLED
#define ASURA_TCP_SOCKET_LOG_ERROR( x ) { if (Asura_TCPSocket::GetLoggingLevel() >= ASURA_TCP_LOG_ERRORS) { Asura_SimpleLogfileOutput(x); } }
#define ASURA_TCP_SOCKET_LOG_EVENT( x ) { if (Asura_TCPSocket::GetLoggingLevel() >= ASURA_TCP_LOG_EVENTS) { Asura_SimpleLogfileOutput(x); } }
#define ASURA_TCP_SOCKET_LOG_EXTRA( x ) { if (Asura_TCPSocket::GetLoggingLevel() >= ASURA_TCP_LOG_EVERYTHING) { Asura_SimpleLogfileOutput(x); } }
#else
#define ASURA_TCP_SOCKET_LOG_ERROR( x )
#define ASURA_TCP_SOCKET_LOG_EVENT( x )
#define ASURA_TCP_SOCKET_LOG_EXTRA( x )
#endif

/// Types /////////////////////////////////////////////////////////////////////////////////////////
// union to avoid casts from SOCKADDR_IN to SOCKADDR which break strict aliasing rules
union Asura_TCP_AddressUnion
{
	SOCKADDR_IN xSockAddrIn;
	SOCKADDR xSockAddr;
};


/// Constants /////////////////////////////////////////////////////////////////////////////////////

// We need the recv buffer to be smaller than the ring buffer so that if data is received which 
// won't fit in what's left of the ring buffer, clearing out complete messages from the ring 
// buffer will allow the newly-received data to be copied in. The required size is the size of 
// the largest message. Consuming all messages will leave at most (max_message_size - 1) bytes in 
// the buffer. 
// NB one byte of the ring buffer is currently unusable due to a limitation in Asura_RingBuffer.
const u_int uASURA_TCP_SOCKET_MAX_MESSAGE_SIZE = 1024;
const u_int uASURA_TCP_SOCKET_RECV_TEMP_BUFFER_SIZE = uASURA_TCP_SOCKET_MAX_MESSAGE_SIZE;
const u_int uASURA_TCP_SOCKET_INCOMING_BUFFER_SIZE = uASURA_TCP_SOCKET_MAX_MESSAGE_SIZE * 2; // ring buffer.

#ifdef ASURA_TCP_LOGGING_ENABLED
const char* const szASURA_TCP_LOGFILE = "logfiles/tcp.txt";
#endif

/// Static Class Data /////////////////////////////////////////////////////////////////////////////

ASURA_TCP_LOGGING_LEVEL Asura_TCPSocket::s_eLoggingLevel = ASURA_TCP_LOG_EVENTS;

/// Helper functions //////////////////////////////////////////////////////////////////////////////

static void InitialiseSockAddr(SOCKADDR_IN& xSockAddr, u_int uAddress, const Asura_NetPort& xPort )
{
	Asura_Mem::Set(&xSockAddr, 0, sizeof(xSockAddr));
	xSockAddr.sin_family = AF_INET;
	xSockAddr.sin_port = htons(xPort.m_usPort);
	xSockAddr.sin_addr.s_addr = uAddress;
#ifdef ASURA_PLATFORM_PS3
	xSockAddr.sin_vport = htons(xPort.m_usVPort);
#endif
}

/// Asura_TCPSocket //////////////////////////////////////////////////////////////////////////////

Asura_TCPSocket::Asura_TCPSocket()
: m_xSocketHandle(INVALID_SOCKET)
{

}

Asura_TCPSocket::~Asura_TCPSocket()
{
	Asura_Assert(m_xSocketHandle == INVALID_SOCKET, "Socket is still bound");
}

bool Asura_TCPSocket::CreateSocket()
{
#ifdef ASURA_PLATFORM_PS3
	m_xSocketHandle = socket(AF_INET, SOCK_STREAM_P2P, IPPROTO_TCP);
#else
	m_xSocketHandle = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
#endif
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Failed to create socket");
	if (m_xSocketHandle >= 0)
	{
		// set socket to non-blocking
#if defined(WINSOCK_VERSION) || defined(_WINSOCKAPI_)
		u_long uNonBlocking = 1;
#ifdef ASURA_ASSERT
		int iResult = 
#endif
			ioctlsocket(m_xSocketHandle, FIONBIO, &uNonBlocking);
		Asura_Assertf( iResult > SOCKET_ERROR, ("Error setting socket to non blocking (0x%x)", WSAGetLastError()));
#else
		int iNonBlocking = 1;
#ifdef ASURA_ASSERT
		int iResult = 
#endif
			setsockopt (m_xSocketHandle, SOL_SOCKET, SO_NBIO, &iNonBlocking, sizeof(int));
		Asura_Assertf( iResult > SOCKET_ERROR, ("Error setting socket to non blocking (0x%x)", WSAGetLastError()));
#endif

#ifdef _XBOX
		{
			// hack to work around sockets not being properly closed.
			BOOL iReuseAddr = 1;
#ifdef ASURA_ASSERT
			int iResult2 = 
#endif
				setsockopt(m_xSocketHandle, SOL_SOCKET, SO_REUSEADDR, reinterpret_cast<const char*>(&iReuseAddr), sizeof(BOOL));
			Asura_Assertf( iResult2 > SOCKET_ERROR, ("Error setting socket to non blocking (0x%x)", WSAGetLastError()));
		}
#endif
		return true;
	}
	return false;
}

bool Asura_TCPSocket::CloseSocket()
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");

	// Is shutdown necessary? - it should only be required if we want to perform graceful 
	// disconnections. In that case we would need to add a state to the socket so that we
	// can wait for it to finish sending data. In practice we can probably just cope with
	// detecting connection loss via error codes and reacting appropriately.
	
#if 0
	const int iShutdownResult = shutdown(m_xSocketHandle, SD_SEND);
	if (iShutdownResult >= 0)
	{
		int iRecvResult = 0;
		do 
		{
			iRecvResult = recv(m_xSocketHandle, NULL, 0, 0);
			Asura_Thread_System::DelayCurrentThread(0.001f);
		}
		while(iRecvResult > 0);
	}
#endif

	const int iCloseResult = closesocket(m_xSocketHandle);	
	Asura_Assertf(iCloseResult >= 0, ("closesocket() returned an error: %i", WSAGetLastError()));

	m_xSocketHandle = INVALID_SOCKET;

	ASURA_TCP_SOCKET_LOG_EVENT((szASURA_TCP_LOGFILE, "Socket closed\n"));

	return (iCloseResult >= 0);
}

#ifdef _DEBUG
const char* Asura_TCPSocket::GetMsgIDDesc(u_int uMsgID)
{
	const char* szDesc = NULL;

#define CASE_TO_STR( _case ) \
		case (_case): (szDesc) = #_case; break

	static char acBuffer[32];
	Asura_PrintfToBfr( acBuffer, "[misc (%d)]", 32, true, uMsgID );
	szDesc = acBuffer;

	switch( uMsgID )
	{
		CASE_TO_STR( ASURA_TCP_MESSAGE_CONNECT_REQUEST );
		CASE_TO_STR( ASURA_TCP_MESSAGE_CONNECT_CONFIRM );
		CASE_TO_STR( ASURA_TCP_MESSAGE_DISCONNECT );

	default:
		//Project_GetMsgIDDesc( uMsgID, pszDesc );
		break;
	}

	return szDesc;

#undef CASE_TO_STR
}
#endif

/// Asura_TCPSocket_Listener /////////////////////////////////////////////////////////////////////

Asura_ReliableConnectionID Asura_TCPSocket_Listener::s_xNextConnectionID = Asura_ReliableConnectionID(uASURA_RELIABLESOCKET_INVALID_CONNECTION_ID + 1);

Asura_TCPSocket_Listener::Asura_TCPSocket_Listener(Asura_TCP_Message_Callback pfnMessageCallback, 
												   Asura_TCP_Connection_Callback pfnConnectionCallback, 
												   Asura_TCP_Disconnection_Callback pfnDisconnectionCallback)
: m_uNumConnections(0)
, m_bActive(false)
{
	for(u_int u = 0; u < uASURA_TCP_SOCKET_MAX_CONNECTIONS; ++u)
	{
		m_apxConnections[u] = new Asura_TCPSocket_Connection(pfnMessageCallback, 
			pfnConnectionCallback, pfnDisconnectionCallback);
	}
}

void Asura_TCPSocket_Listener::SetCallbackUserData( void* pUserData )
{
	for(u_int u = 0; u < uASURA_TCP_SOCKET_MAX_CONNECTIONS; ++u)
	{
		m_apxConnections[u]->SetCallbackUserData(pUserData);
	}
}

Asura_TCPSocket_Listener::~Asura_TCPSocket_Listener()
{
	for(u_int u = 0; u < uASURA_TCP_SOCKET_MAX_CONNECTIONS; ++u)
	{
		delete m_apxConnections[u];
	}
}

bool Asura_TCPSocket_Listener::Listen(const Asura_NetPort& xPort)
{
	return Listen(INADDR_ANY, xPort);
}

bool Asura_TCPSocket_Listener::Listen(u_int uAddress, const Asura_NetPort& xPort)
{
	Asura_Assert(!IsActive(), "already listening");
	CreateSocket();

	Asura_TCP_AddressUnion xSockAddr;
	InitialiseSockAddr(xSockAddr.xSockAddrIn, uAddress, xPort);

	const int iBindResult = bind(m_xSocketHandle, &xSockAddr.xSockAddr, sizeof(xSockAddr));
	if(0 > iBindResult)
	{
#if defined (ASURA_TCP_LOGGING_ENABLED) || defined (ASURA_ASSERT)
		const int iError = WSAGetLastError();
#endif
		ASURA_TCP_SOCKET_LOG_ERROR((szASURA_TCP_LOGFILE, "Listen failed: bind error 0x%08x\n", iError));
		Asura_Assertf(0, ("Failed to bind socket: %i", iError));
		CloseSocket();
		return false;
	}

	const int iListenResult = listen(m_xSocketHandle, uASURA_TCP_SOCKET_MAX_CONNECTIONS);

	if(0 > iListenResult)
	{
#if defined (ASURA_TCP_LOGGING_ENABLED) || defined (ASURA_ASSERT)
		const int iError = WSAGetLastError();
#endif
		ASURA_TCP_SOCKET_LOG_ERROR((szASURA_TCP_LOGFILE, "Listen failed: 0x%08x\n", iError));
		Asura_Assertf(0, ("Failed to listen on socket: %i", iError));
		CloseSocket();
		return false;
	}

	ASURA_TCP_SOCKET_LOG_EVENT((szASURA_TCP_LOGFILE, "Socket listening.\n"));

	m_bActive = true;
	return true;
}

bool Asura_TCPSocket_Listener::SendToAll(u_short usMsgID, u_short usMsgBodySize, const void* pMsgBody)
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	bool bSuccess = true;
	for(u_int u = 0; u < m_uNumConnections; ++u)
	{
		Asura_TCPSocket_Connection* pxConnection = m_apxConnections[u];
		if(pxConnection && pxConnection->IsConnected())
		{
			bSuccess &= pxConnection->SendMessage(usMsgID, usMsgBodySize, pMsgBody);
			Asura_Assertf(bSuccess, ("Failed to send data on connection ID %i", pxConnection->GetIDAsInt()));
		}
	}
	return bSuccess;
}

bool Asura_TCPSocket_Listener::SendToAllExcept(u_short usMsgID, u_short usMsgBodySize, 
											   const void* pMsgBody, Asura_ReliableConnectionID xConnectionID)
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	bool bSuccess = true;
	for(u_int u = 0; u < m_uNumConnections; ++u)
	{
		Asura_TCPSocket_Connection* pxConnection = m_apxConnections[u];
		if(pxConnection && pxConnection->IsConnected() && (pxConnection->GetID() != xConnectionID))
		{
			bSuccess &= pxConnection->SendMessage(usMsgID, usMsgBodySize, pMsgBody);
			Asura_Assertf(bSuccess, ("Failed to send data on connection ID %i", pxConnection->GetIDAsInt()));
		}
	}
	return bSuccess;
}

bool Asura_TCPSocket_Listener::SendMessage(Asura_ReliableConnectionID xConnectionID, u_short usMsgID, u_short usMsgBodySize, const void* pMsgBody)
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	// Find the relevant connection and send on that.
	Asura_TCPSocket_Connection* pxConnection = GetConnection(xConnectionID);
	Asura_Assert(pxConnection, "No connection for given address");
	if(pxConnection)
	{
		return pxConnection->SendMessage(usMsgID, usMsgBodySize, pMsgBody);
	}
	return false;
}

Asura_TCPSocket_Connection* Asura_TCPSocket_Listener::GetConnection(Asura_ReliableConnectionID xID)
{
	// Flat search should be sufficient as long as we're not doing this too often.
	for(u_int u = 0; u < m_uNumConnections; ++u)
	{
		Asura_TCPSocket_Connection* pxConnection = m_apxConnections[u];
		Asura_Assert(pxConnection, "No connection");
		if(pxConnection && pxConnection->GetID() == xID)
		{			
			return pxConnection;
		}
	}
	return NULL;
}

void Asura_TCPSocket_Listener::Update()
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	fd_set xFDSet;
	FD_ZERO(&xFDSet);
	FD_SET(m_xSocketHandle, &xFDSet);
	struct timeval timeout = {0, 0};

	// Listen for new connections
	int iSelectResult = 0;
	do
	{
		iSelectResult = select(FD_SETSIZE, &xFDSet, NULL, NULL, &timeout);
		if (iSelectResult > 0)
		{
			if(m_uNumConnections < uASURA_TCP_SOCKET_MAX_CONNECTIONS)
			{
				// Find an unused connection
				Asura_TCPSocket_Connection* pxConnection = m_apxConnections[m_uNumConnections];
				Asura_Assert(pxConnection, "No connection");
				if(pxConnection && pxConnection->Accept(*this, s_xNextConnectionID))
				{
					++m_uNumConnections;
				}
				s_xNextConnectionID++;
			}
			else
			{
				// Too many connections. Accept and close immediately
				SOCKET xSocketHandle = accept(m_xSocketHandle, NULL, NULL);
				// TODO - Send a close reason first. It won't be fully guaranteed
				// due to the close, but this should be a rare situation anyway.
				closesocket(xSocketHandle);

				ASURA_TCP_SOCKET_LOG_ERROR((szASURA_TCP_LOGFILE, "rejected new connection: already maxed out\n"));
			}
		}
	} while (iSelectResult > 0);

#ifdef ASURA_ASSERT
	if (iSelectResult < 0)
	{
		int iLastError = WSAGetLastError();
		Asura_Assertf(false, ("TCP Listener: Error %d calling select", iLastError));
	}
#endif

	// Update existing connections in reverse so that when we remove one
	// we don't need to worry about fixing up the loop counter.
	u_int u = m_uNumConnections;
	while(u--)
	{
		Asura_TCPSocket_Connection* pxConnection = m_apxConnections[u];
		Asura_Assert(pxConnection, "No connection");
		if(pxConnection)
		{
			pxConnection->Update();
			if(!pxConnection->IsActive())
			{
				// Connection has closed. Remove it.
				ASURA_TCP_SOCKET_LOG_EVENT((szASURA_TCP_LOGFILE, "closing connection %u (no longer active)\n", pxConnection->GetIDAsInt()));

				// Swap the connections about so we always have all valid 
				// connections at the start of the array.
				--m_uNumConnections;
				Asura_TCPSocket_Connection *pxTemp = m_apxConnections[m_uNumConnections];
				m_apxConnections[m_uNumConnections] = m_apxConnections[u];
				m_apxConnections[u] = pxTemp;
			}
		}
	}
}

void Asura_TCPSocket_Listener::Close()
{
	Asura_Assert(IsActive(), "Close() called on inactive socket");

	// Close all connections
	while(m_uNumConnections)
	{
		m_apxConnections[--m_uNumConnections]->Disconnect(true, ASURA_TCP_DISCONNECT_SHUTDOWN);
	}
	CloseSocket();

	m_bActive = false;
}

void Asura_TCPSocket_Listener::Disconnect(Asura_ReliableConnectionID xConnectionID, ASURA_TCP_DISCONNECT_REASON eReason)
{
	Asura_TCPSocket_Connection* pxConnection = GetConnection(xConnectionID);
	Asura_Assert(pxConnection, "No connection with that ID");
	if(pxConnection)
	{
		ASURA_TCP_SOCKET_LOG_EVENT((szASURA_TCP_LOGFILE, "Disconnect connection %u, reason %u\n", xConnectionID.AsInt(), eReason));
		pxConnection->Disconnect(false, eReason);
	}
}

u_short Asura_TCPSocket_Listener::GetListenPort() const
{
	if (m_xSocketHandle == INVALID_SOCKET)
	{
		return 0;
	}
	Asura_TCP_AddressUnion xSockAddr;
	int iAddressSize = sizeof(xSockAddr);
	int iRet = getsockname(m_xSocketHandle, &xSockAddr.xSockAddr, &iAddressSize);
	if (iRet < 0)
	{
		Asura_Assertf(false, ("getsockname() failed (0x%x)", WSAGetLastError()));
		return 0;
	}
	return htons(xSockAddr.xSockAddrIn.sin_port);
}

/// Asura_TCPSocket_Connection ///////////////////////////////////////////////////////////////////

Asura_TCPSocket_Connection::Asura_TCPSocket_Connection(Asura_TCP_Message_Callback pfnCallback, 
													   Asura_TCP_Connection_Callback pfnConnectionCallback, 
													   Asura_TCP_Disconnection_Callback pfnDisconnectionCallback)
: m_eState(STATE_DISCONNECTED)
, m_bWaitingForMessageBody(false)
, m_usConnectionDataSize(0)
, m_pConnectionData(NULL)
, m_pfnMessageCallback(pfnCallback)
, m_pfnConnectionCallback(pfnConnectionCallback)
, m_pfnDisconnectionCallback(pfnDisconnectionCallback)
, m_pCallbackUserData(NULL)
{
	m_xIncomingDataBuffer.Create(uASURA_TCP_SOCKET_INCOMING_BUFFER_SIZE);
	Asura_Mem::Set(&m_xCurrentMessageHeader, 0, sizeof(m_xCurrentMessageHeader));
}

Asura_TCPSocket_Connection::~Asura_TCPSocket_Connection()
{
	m_xIncomingDataBuffer.Destroy();
	Asura_Deallocate(m_pConnectionData);
}

void Asura_TCPSocket_Connection::Update()
{
	Asura_Assert(IsActive(), "Should not be updating connection if not connected");

	if (m_eState == STATE_CONNECTING)
	{
		// wait for connection
		struct timeval xTimeout = {0, 0};
		fd_set xWriteFDSet;
		fd_set xExceptFDSet;
		FD_ZERO(&xWriteFDSet);
		FD_SET(m_xSocketHandle, &xWriteFDSet);
		FD_ZERO(&xExceptFDSet);
		FD_SET(m_xSocketHandle, &xExceptFDSet);

		int iRet = select(FD_SETSIZE, NULL, &xWriteFDSet, &xExceptFDSet, &xTimeout);
#ifdef ASURA_ASSERT
		if (iRet < 0)
		{
			int iLastError = WSAGetLastError();
			Asura_Assertf( false, ("TCP Connection: Error %d calling select", iLastError ));
		}
#endif

		if (iRet > 0)
		{
			if (FD_ISSET(m_xSocketHandle, &xExceptFDSet))
			{
				// connect failed
				Asura_Deallocate(m_pConnectionData);
				m_pConnectionData = 0;
				m_usConnectionDataSize = 0;
				Disconnect(true, ASURA_TCP_DISCONNECT_CONNECTION_TIMEOUT);
			}
			else if (FD_ISSET(m_xSocketHandle, &xWriteFDSet))
			{
				// connect succeeded, send initial message
				m_eState = STATE_REQUESTING;
				Asura_Assert(m_usConnectionDataSize && m_pConnectionData, "no connection data?");
				SendMessage(ASURA_TCP_MESSAGE_CONNECT_REQUEST, m_usConnectionDataSize, m_pConnectionData);
				// don't need the initial message anymore
				Asura_Deallocate(m_pConnectionData);
				m_pConnectionData = 0;
				m_usConnectionDataSize = 0;
			}
		}
		return;
	}

	UpdateIncomingData();
	// Now handle any messages in the buffer.
	// (MH - 21.04.2009) see comment in UpdateIncomingData() where this function is also called. This one
	// is in theory redundant as a result.
	HandleMessages();
}

void Asura_TCPSocket_Connection::UpdateIncomingData()
{
	struct timeval xTimeout = {0, 0};
	fd_set xFDSet;
	FD_ZERO(&xFDSet);
	FD_SET(m_xSocketHandle, &xFDSet);
	// Check for incoming data
	char acTempBuffer[uASURA_TCP_SOCKET_RECV_TEMP_BUFFER_SIZE];
	while(select(FD_SETSIZE, &xFDSet, NULL, NULL, &xTimeout) > 0)
	{
		// (MH - 21.04.2009) reading in message header then message body repeatedly and handling the messages
		// seems to work around an apparent bug in the ring buffer.
		// TODO: figure out what is going wrong and remove this? Maybe it's not a big deal to just read (header, data,
		// header, data, ...) with multiple recv calls anyway...
		u_int uMaxToReceive = m_bWaitingForMessageBody ? uASURA_TCP_SOCKET_RECV_TEMP_BUFFER_SIZE : sizeof(Asura_TCP_MessageHeader);

		const int iBytesReceived = recv(m_xSocketHandle, acTempBuffer, uMaxToReceive, 0);
		if(iBytesReceived > static_cast<int>(m_xIncomingDataBuffer.GetUnusedSpace()))
		{
			// Too much data in buffer. Attempt to clear it out.
			HandleMessages();
			if(iBytesReceived > static_cast<int>(m_xIncomingDataBuffer.GetUnusedSpace()))
			{
				// Hmmm, still too much. Give up.
				Asura_Assert( false, "Disconnecting due to buffer overflow" ); 
				Disconnect(false, ASURA_TCP_DISCONNECT_BUFFER_FULL);
			}			
		}
		// Zero denotes connection has been gracefully closed.
		else if(iBytesReceived == 0)
		{
			Asura_Assert(m_eState == STATE_DISCONNECTING, "Was not expecting connection to close");
			Disconnect(true, ASURA_TCP_DISCONNECT_SHUTDOWN);
			break;
		}
		else if(iBytesReceived == SOCKET_ERROR)
		{
			// Connection has been closed
			int iLastError = WSAGetLastError();
			// The 360 is returning WSAECONNRESET on a graceful connection. As its 
			// never an error as such, treat it as a normal shutdown. PJC 30/04/09
			if (iLastError == WSAECONNRESET)
			{
				Disconnect(true, ASURA_TCP_DISCONNECT_SHUTDOWN);
			}
			else
			{
				Asura_Assertf(false, ("TCP connection closed with error: %i", iLastError));
				Disconnect(true, ASURA_TCP_DISCONNECT_SOCKET_ERROR);
			}
			break;
		}
		else
		{
			Asura_Assertf(iBytesReceived > 0, ("Bad return value from recv(): %i", iBytesReceived));
#ifdef ASURA_ASSERT
			const u_int uBytesWritten = 
#endif
				m_xIncomingDataBuffer.Write(acTempBuffer, iBytesReceived);
			Asura_Assert(static_cast<int>(uBytesWritten) == iBytesReceived, "Did not write expected number of bytes");
			
			// (MH - 21.04.2009) handle messages every time we receive some data
			HandleMessages();
		}
	}
}

void Asura_TCPSocket_Connection::HandleMessages()
{
	while(true)
	{
		const u_int uDataAvailable = m_xIncomingDataBuffer.GetCount();
		if(!uDataAvailable)
		{
			// Nothing to parse
			break;
		}
		if(m_bWaitingForMessageBody)
		{
			const u_short usDataSize = m_xCurrentMessageHeader.m_usDataSize;
			if(uDataAvailable >= usDataSize)
			{
				// Could just create this on the stack with maximum message size?
				char* pcMessageData = new char[usDataSize];
				// Read in the message body
				m_xIncomingDataBuffer.Read(pcMessageData, usDataSize);
				// Check the checksum.
				const u_int uCheckSum = Asura_GetHashID_Binary(pcMessageData, usDataSize);
				// testchecksum: Asura_OutputToDebugger(("checksum out (%u bytes): 0x%x\n", usDataSize, uCheckSum));
				Asura_Assert(uCheckSum == m_xCurrentMessageHeader.m_uCheckSum, "Checksum does not match");
				// Only process the message if the checksum matches
				if(uCheckSum == m_xCurrentMessageHeader.m_uCheckSum)
				{
					// Handle the message
					HandleMessage(m_xCurrentMessageHeader.m_usMessageID, usDataSize, pcMessageData);
				}
				else
				{
					Disconnect(true, ASURA_TCP_DISCONNECT_MESSAGE_CHECKSUM_ERROR);
				}
				m_bWaitingForMessageBody = false;
				delete [] pcMessageData;
			}
			else
			{
				// Body is incomplete - we're done for this update.
				break;
			}
		}
		else
		{
			// See if there's a header's worth available.
			if(uDataAvailable >= sizeof(Asura_TCP_MessageHeader))
			{
				m_xIncomingDataBuffer.Read(reinterpret_cast<char*>(&m_xCurrentMessageHeader), sizeof(Asura_TCP_MessageHeader));
				if(m_xCurrentMessageHeader.m_usDataSize)
				{
					m_bWaitingForMessageBody = true;
				}
				else
				{
					// Handle it right now, as we don't need any further data.
					HandleMessage(m_xCurrentMessageHeader.m_usMessageID, 0, NULL);
				}
			}
			else
			{
				// Header is incomplete - we're done for this update.
				break;
			}
		}
	}
}

void Asura_TCPSocket_Connection::HandleMessage(u_short usMsgID, u_short usMsgBodySize, const char* pcMsgBody)
{
	// Check whether it's a connection control message.
	if(usMsgID >= ASURA_TCP_MESSAGE_FIRST_RESERVED_ID)
	{
		switch(usMsgID)
		{
			case ASURA_TCP_MESSAGE_CONNECT_REQUEST:
			{
				Asura_Assert(STATE_AWAITING_REQUEST == m_eState, "Unexpected state");
				Asura_Assert(m_pfnConnectionCallback, "No connection callback");
				if(m_pfnConnectionCallback && m_pfnConnectionCallback(m_xID, usMsgBodySize, pcMsgBody, m_pCallbackUserData))
				{
					// Connection approved. Send a confirm back.
					SendMessage(ASURA_TCP_MESSAGE_CONNECT_CONFIRM, 0, NULL);
					m_eState = STATE_CONNECTED;
				}
				else
				{
					// Send out a disconnection message.
					Disconnect(false, ASURA_TCP_DISCONNECT_CONNECTION_REJECTED);
				}

				break;
			}
			case ASURA_TCP_MESSAGE_CONNECT_CONFIRM:
			{
				Asura_Assert(STATE_REQUESTING == m_eState, "Unexpected state");
				m_eState = STATE_CONNECTED;
				break;
			}
			case ASURA_TCP_MESSAGE_DISCONNECT:
			{
				const ASURA_TCP_DISCONNECT_REASON* peReason = Asura_ReliableSocket_CheckAndExtractData<ASURA_TCP_DISCONNECT_REASON>( pcMsgBody, usMsgBodySize);
				Disconnect(true, (peReason) ? *peReason : ASURA_TCP_DISCONNECT_UNKNOWN );
				break;
			}
		}
	}
	else
	{
		//Don't assert about non-reserved messages received while disconnecting, as we have no control over them. Just ignore them.
		Asura_Assert(IsConnected() || (m_eState == STATE_DISCONNECTING),"Trying to handle a message when not connected");
		if (IsConnected())
		{
			Asura_Assert(m_pfnMessageCallback, "No message callback");
			if(m_pfnMessageCallback)
			{
				m_pfnMessageCallback(m_xID, usMsgID, usMsgBodySize, pcMsgBody, m_pCallbackUserData);
			}
		}
	}
}

bool Asura_TCPSocket_Connection::SendMessage(u_short usMsgID, u_short usMsgBodySize, const void* pMsgBody)
{
	Asura_Assert(usMsgBodySize <= uASURA_TCP_SOCKET_MAX_MESSAGE_SIZE, "message too big, closing connection");
	if (usMsgBodySize > uASURA_TCP_SOCKET_MAX_MESSAGE_SIZE)
	{
		Disconnect(true, ASURA_TCP_DISCONNECT_MESSAGE_TOO_BIG);
		return false;
	}

	Asura_Assert(IsConnected() || (usMsgID >= ASURA_TCP_MESSAGE_FIRST_RESERVED_ID),"Trying to send data when not connected");
	if (IsConnected() || (usMsgID >= ASURA_TCP_MESSAGE_FIRST_RESERVED_ID))
	{
		Asura_TCP_MessageHeader xHeader;
		xHeader.m_usMessageID = usMsgID;
		xHeader.m_usDataSize = usMsgBodySize;
		xHeader.m_uCheckSum = Asura_GetHashID_Binary(pMsgBody, usMsgBodySize);
		// testchecksum: Asura_OutputToDebugger(("checksum in (%u bytes): 0x%x\n", usMsgBodySize, xHeader.m_uCheckSum));
		if(Send(&xHeader, sizeof(xHeader)))
		{
			return usMsgBodySize ? Send(pMsgBody, usMsgBodySize) : true;
		}
	}
	return false;
}

bool Asura_TCPSocket_Connection::Connect(u_int uAddress, const Asura_NetPort& xPort, u_short usConnectionDataSize, const void* pConnectionData)
{
	CreateSocket();

	Asura_TCP_AddressUnion xSockAddr;
	InitialiseSockAddr(xSockAddr.xSockAddrIn, uAddress, xPort);
	const int iConnectResult = connect(m_xSocketHandle, &xSockAddr.xSockAddr, sizeof(xSockAddr));
	const int iLastError = WSAGetLastError();
	// on PC a non-blocking socket returns WSAEWOULDBLOCK. On PS3 it returns WSAEINPROGRESS.
	// Therefore we check both. Note that on PS3 WSAEWOULDBLOCK represents a timeout on a blocking socket.
	Asura_Assertf(iConnectResult >= 0 || iLastError == WSAEWOULDBLOCK || iLastError == WSAEINPROGRESS, 
		("connect() returned an error: %i", iLastError ));
	if (iConnectResult < 0 && iLastError != WSAEWOULDBLOCK && iLastError != WSAEINPROGRESS)
	{
		ASURA_TCP_SOCKET_LOG_ERROR((szASURA_TCP_LOGFILE, "connect() failed: 0x%08x\n", iLastError));
		Disconnect(true, ASURA_TCP_DISCONNECT_SOCKET_ERROR);
		return false;
	}
	else
	{
		m_xID.Invalidate();
		m_eState = STATE_CONNECTING;
		// allocate the buffer to store the connection data so we can use it when the connect has completed.
		m_usConnectionDataSize = usConnectionDataSize;
		m_pConnectionData = Asura_Allocate(usConnectionDataSize);
		Asura_Mem::Copy(m_pConnectionData, pConnectionData, usConnectionDataSize);
		// ensure the the incoming data buffer and header flag are initialised.
		m_xIncomingDataBuffer.Reset();
		m_bWaitingForMessageBody = false;
		return true;
	}
}

bool Asura_TCPSocket_Connection::Accept(class Asura_TCPSocket_Listener& xListenerSocket, Asura_ReliableConnectionID xID)
{
	m_xID = xID;
	m_xSocketHandle = accept(xListenerSocket.GetSocket(), NULL, NULL);
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	m_eState = STATE_AWAITING_REQUEST;
	bool bOK = (m_xSocketHandle != INVALID_SOCKET);
	if (bOK)
	{
		ASURA_TCP_SOCKET_LOG_EVENT((szASURA_TCP_LOGFILE, "accepted new connection: %u\n", xID.AsInt()));
	}
	else
	{
		ASURA_TCP_SOCKET_LOG_ERROR((szASURA_TCP_LOGFILE, "rejected connection: accept failed\n"));
	}
	return bOK;
}

bool Asura_TCPSocket_Connection::Send(const void* pData, u_int uBytes)
{
	Asura_Assert(m_xSocketHandle != INVALID_SOCKET, "Invalid socket");
	const int iBytesSent = send(m_xSocketHandle, reinterpret_cast<const char*>(pData), uBytes, 0);
	Asura_Assertf(iBytesSent > 0, ("Failed to send data. Error code: %i", WSAGetLastError()));
	return (iBytesSent == static_cast<int>(uBytes));
}

bool Asura_TCPSocket_Connection::GetIncomingData(char* pcBuffer, u_int uNumBytes)
{
	Asura_Assert(uNumBytes <= m_xIncomingDataBuffer.GetCount(), "Requested more data than we have available");
	const u_int uBytesRead = m_xIncomingDataBuffer.Read(pcBuffer, uNumBytes);
	return (uBytesRead == uNumBytes);
}

void Asura_TCPSocket_Connection::Disconnect(bool bImmediate, ASURA_TCP_DISCONNECT_REASON eReason)
{
	Asura_Assert(m_eState != STATE_DISCONNECTED,"Trying to disconnect or start disconnecting when already disconnected");
	if (m_eState != STATE_DISCONNECTED)
	{
		if(bImmediate)
		{
			// ensure the the incoming data buffer and header flag are reset.
			m_xIncomingDataBuffer.Reset();
			m_bWaitingForMessageBody = false;

			CloseSocket();
			m_eState = STATE_DISCONNECTED;

			if (m_pfnDisconnectionCallback)
			{
				m_pfnDisconnectionCallback(m_xID, eReason, m_pCallbackUserData);
			}
		}
		else
		{
			Asura_Assert(m_eState != STATE_DISCONNECTING, "Trying to start disconnecting while already disconnecting");
			if (m_eState != STATE_DISCONNECTING)
			{
				SendMessage(ASURA_TCP_MESSAGE_DISCONNECT, sizeof(ASURA_TCP_DISCONNECT_REASON), reinterpret_cast<const char*>(&eReason));
				m_eState = STATE_DISCONNECTING;
			}
		}
	}
}

#endif // #ifdef ASURA_SUPPORT_EXTERNAL_NETWORK
