#include "pch.h"
#include "ClientSocket.h"
#include <mmsystem.h>

// Constructor
CClientSocket::CClientSocket()
	: m_WsaData()
	, m_Address()
	, m_Socket()
{
}

// Destructor
CClientSocket::~CClientSocket()
{
}

// Open
int CClientSocket::Connect(const char *pszIpAddress, unsigned int Port)
{
	unsigned long Argp = 1;
	char OptVal = 1;
	int RetVal;
	int ErrCode;

	// Initialize WinSock
	RetVal = ::WSAStartup(MAKEWORD(2, 0), &m_WsaData);
	if (RetVal != 0){
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	// Create socket
	m_Socket = ::socket(AF_INET, SOCK_STREAM, 0);
	if (m_Socket == INVALID_SOCKET) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	// Convert address
	m_Address.sin_family = AF_INET;
	m_Address.sin_port = htons(static_cast<unsigned short>(Port));
	::inet_pton(AF_INET, pszIpAddress, &m_Address.sin_addr.s_addr);

	// Connect
	RetVal = ::connect(m_Socket, reinterpret_cast<sockaddr *>(&m_Address), sizeof(m_Address));
	if (RetVal != 0) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	// Set to non-blocking
	RetVal = ::ioctlsocket(m_Socket, FIONBIO, &Argp);
	if (RetVal != 0) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	// Disable transmit delay (Nagle algorithm)
	RetVal = ::setsockopt(m_Socket, IPPROTO_TCP, TCP_NODELAY, &OptVal, sizeof(OptVal));
	if (RetVal != 0) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	return 0;
}

// Close
int CClientSocket::Disconnect()
{
	int RetVal;
	int ErrCode;

	// Close
	RetVal = ::closesocket(m_Socket);
	if (RetVal != 0) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	// Terminate WinSock
	RetVal = ::WSACleanup();
	if (RetVal != 0) {
		ErrCode = ::GetLastError();
		return ErrCode;
	}

	return 0;
}

// Transmit
int CClientSocket::Send(const char *pszBuffer)
{
	int Offset;
	int Remain;
	int RetVal;
	int ErrCode;

	// Write data to socket
	Offset = 0;
	Remain = static_cast<int>(::strlen(pszBuffer));
	while (Remain > 0) {
		RetVal = ::send(m_Socket, pszBuffer + Offset, Remain, 0);
		if (RetVal == SOCKET_ERROR) {
			ErrCode = ::GetLastError();
			return ErrCode;
		}
		Offset += RetVal;
		Remain -= RetVal;
	}

	return 0;
}

// Receive
int CClientSocket::Receive(CString &strBuffer, unsigned long Timeout)
{
	DWORD StartTime;
	int RetVal;
	int ErrCode;

	// Read data from the socket
	strBuffer = _T("");
	StartTime = ::timeGetTime();
	while (1) {
		::memset(m_szBuffer, 0, sizeof(m_szBuffer));
		RetVal = ::recv(m_Socket, m_szBuffer, sizeof(m_szBuffer) - 1, 0);
		if (RetVal == SOCKET_ERROR) {
			if (::WSAGetLastError() != WSAEWOULDBLOCK) {
				ErrCode = ::WSAGetLastError();
				return ErrCode;
			}
			break;
		}
		else if (RetVal == 0) {
			break;
		}
		else {
			m_szBuffer[RetVal] = '\0';
			CString strFragment(m_szBuffer);
			strBuffer = strBuffer + strFragment;
		}
		if (::timeGetTime() - StartTime >= Timeout) {
			ErrCode = WSAETIMEDOUT;
			return ErrCode;
		}
	}

	return 0;
}
