#include "Winsock2.h"
#include "Windows.h"
#include "stdio.h"
#include "TCPSocket.h"

TCPSocket::TCPSocket()
{
	memset((void *)this, 0, sizeof(TCPSocket));
}

TCPSocket::TCPSocket(SOCKET &Socket, char Type)
{
	memset((void *)this, 0, sizeof(TCPSocket));
	this->m_Socket = Socket;
	this->m_Mode = Type;
	this->m_Connected = 1;

	m_Event = WSACreateEvent();
	WSAEventSelect(this->m_Socket, m_Event, FD_READ);
}

TCPSocket::TCPSocket(const char *Host, int Port)
{
	memset((void *)this, 0, sizeof(TCPSocket));
	this->Client(Host, Port);
}

TCPSocket::TCPSocket(const char *IP, int Port, int Backlog)
{
	memset((void *)this, 0, sizeof(TCPSocket));
	this->Server(IP, Port, Backlog);
}

TCPSocket::TCPSocket(TCPSocket &Socket)
{
	memcpy((void *)this, (void *)&Socket, sizeof(TCPSocket));
}


TCPSocket::~TCPSocket()
{
	this->Destroy();
}


bool TCPSocket::Startup()
{
	WORD wVersionRequested = MAKEWORD(2, 2);
	WSADATA wsaData;
	int err = WSAStartup(wVersionRequested, &wsaData);
	if (err != 0)
	{
		printf("WSAStartup failed with error: %d\n", err);
		return 0;
	}

	if (LOBYTE(wsaData.wVersion) != 2 || HIBYTE(wsaData.wVersion != 2))
	{
		printf("You need Winsock 2.2.\n");
		WSACleanup();
		return 0;
	}
	return 1;
}

bool TCPSocket::Shutdown()
{
	return WSACleanup() == 0 ? 1 : 0;
}


bool TCPSocket::Client(const char *Host, int Port)
{
	hostent *he;
	if ((he = gethostbyname(Host)) == 0) 
	{
		return 0;
	}
	SOCKET s_ = socket(AF_INET,SOCK_STREAM,0);
	if (s_ == INVALID_SOCKET) 
	{
		return 0;
	}
	sockaddr_in addr;
	addr.sin_family = AF_INET;
	addr.sin_port = htons((u_short)Port);
	addr.sin_addr = *((in_addr *)he->h_addr);
	memset(&(addr.sin_zero), 0, 8); 
	if (connect(s_, (sockaddr *) &addr, sizeof(sockaddr))) 
	{
		return 0;
	}

	this->m_Socket = s_;
	this->m_Connected = 1;
	this->m_Mode = 0;

	m_Event = WSACreateEvent();
	WSAEventSelect(this->m_Socket, m_Event, FD_READ);
	return 1;
}

bool TCPSocket::Server(const char *IP, int Port, int Backlog)
{
	SOCKET server = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (server == INVALID_SOCKET) 
	{
		return 0;
	}

	sockaddr_in service;
	service.sin_family = AF_INET;
	service.sin_addr.s_addr = inet_addr(IP == 0 ? "127.0.0.1" : IP);
	service.sin_port = htons((u_short)Port);

	if(bind(server, (SOCKADDR*) &service, sizeof(service)) == SOCKET_ERROR) 
	{
		return 0;
	}

	if (listen(server, Backlog) == SOCKET_ERROR) 
	{
		return 0;
	}
	
	this->m_Socket = server;
	this->m_Mode = 1;
	this->m_Connected = 1;

	m_Event = WSACreateEvent();
	WSAEventSelect(this->m_Socket, m_Event, FD_ACCEPT);

	return 1;
}

bool TCPSocket::Sendf(const char *Format, ...)
{
	char buffer[1024];
	va_list va;
	_crt_va_start(va, Format);
	vsnprintf_s(buffer, 256, 256, Format, va);
	va_end(va);

	if(!this->Is_Connected())
	{
		return 0;
	}
	int Length = (int)strlen(buffer);
	int Send_ = send(this->m_Socket, buffer, Length, 0);
	if(Send_ != Length)
	{
		if(Send_ == -1)
		{
			return 0;
		}
		else
		{
			this->Send(buffer+Send_, Length-Send_);
		}
	}

	return 1;
}

bool TCPSocket::Send(const char *Data, int Length)
{
	if(!this->Is_Connected())
	{
		return 0;
	}
	int Send_ = send(this->m_Socket, Data, Length, 0);
	if(Send_ != Length)
	{
		if(Send_ == -1)
		{
			return 0;
		}
		else
		{
			this->Send(Data+Send_, Length-Send_);
		}
	}

	return 1;
}

bool TCPSocket::Receive(char *Buffer, int Length)
{
	if(this->Is_DataAvaliable())
	{
		int ret = recv(this->m_Socket, Buffer, Length, 0);
		if(ret <= 0)
		{
			this->Destroy();
			return 0;
		}
	}
	else
	{
		return 0;
	}
	return 1;
}

bool TCPSocket::Accept(TCPSocket *Socket)
{
	if(this->Is_DataAvaliable())
	{
		sockaddr_in addr;
		int s_sa = sizeof(sockaddr);
		SOCKET client = accept(this->m_Socket, (sockaddr *)&addr, &s_sa);
		if(client == INVALID_SOCKET) 
		{
			return 0;
		}

		Socket->TCPSocket::TCPSocket(client, 1);
		return 1;
	}
	return 0;
}

bool TCPSocket::Destroy()
{
	shutdown(this->m_Socket, SD_BOTH);
	closesocket(this->m_Socket);
	WSACloseEvent(this->m_Event);

	memset((void *)this, 0, sizeof(TCPSocket));
	return 0;
}

bool TCPSocket::Is_Connected()
{
	return m_Connected;
}

bool TCPSocket::Is_DataAvaliable()
{
	if(!m_Connected)
	{
		return 0;
	}
	DWORD ret = WaitForSingleObject(m_Event, 0);
	if(ret == WAIT_FAILED)
	{
		return 0;
	}
	return ret == WAIT_TIMEOUT ? 0 : 1;
}

