#include <WinSock2.h>
#include <WS2tcpip.h>
#pragma comment(lib,"ws2_32.lib")
#include <iostream>

#define HOSTIP "192.168.0.11"

enum RunType
{
	Server,
	Client
};

RunType runtype = RunType::Server;

void DoServer()
{
	SOCKET serverSocketHandle = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

	if (serverSocketHandle == NULL)
	{
		std::cerr << "Failed to create socket" << std::endl;
		return;
	}

	DWORD nonBlocking = 1;
	if (ioctlsocket(serverSocketHandle,	FIONBIO, &nonBlocking) != 0)
	{
		std::cerr << "Failed to set socket non-blocking." << std::endl;
		return;
	}

	sockaddr_in listenAddress;
	listenAddress.sin_family = AF_INET;
	listenAddress.sin_addr.s_addr = INADDR_ANY;
	listenAddress.sin_port = htons(8000);
	int listenAddressSize = sizeof(listenAddress);

	if (bind(serverSocketHandle,(const sockaddr*)&listenAddress,sizeof(sockaddr_in)) < 0)
	{
		std::cerr << "Failed to bind socket." << std::endl;
		return;
	}

	uint32_t recvPacketCounter = 0;

	while (true)
	{
		if (GetAsyncKeyState(VK_ESCAPE))
			return;
		if (recvfrom(serverSocketHandle, (char*)(&recvPacketCounter), sizeof(uint32_t), NULL, (sockaddr*)&listenAddress, &listenAddressSize) > 0) //Server is receiving messages
		{
			sendto(serverSocketHandle, (char*)&recvPacketCounter, sizeof(uint32_t), NULL, (sockaddr*)&listenAddress, sizeof(sockaddr_in)); //This send is never being picked up by the client?
			recvPacketCounter = ntohl(recvPacketCounter);
			std::cout << "Server - [Sender Port: " << ntohs(listenAddress.sin_port) << "] Echoing packet: " << recvPacketCounter << std::endl;
		}
	}

}

void DoClient()
{
	SOCKET clientSocketHandle = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);

	if (clientSocketHandle == NULL)
	{
		std::cerr << "Failed to create socket" << std::endl;
		return;
	}

	DWORD nonBlocking = 1;
	if (ioctlsocket(clientSocketHandle, FIONBIO, &nonBlocking) != 0)
	{
		std::cerr << "Failed to set socket non-blocking." << std::endl;
		return;
	}

	uint32_t ip = INADDR_NONE;
	inet_pton(AF_INET, HOSTIP, &ip);
	if (ip == INADDR_NONE)
	{
		std::cerr << "Failed to resolve host." << std::endl;
		return;
	}

	uint32_t hostip = INADDR_NONE;
	inet_pton(AF_INET, HOSTIP, &hostip);
	if (hostip == INADDR_NONE)
	{
		std::cerr << "Failed to resolve host ip." << std::endl;
		return;
	}

	sockaddr_in sendtoAddress;
	sendtoAddress.sin_family = AF_INET;
	sendtoAddress.sin_addr.s_addr = hostip;
	sendtoAddress.sin_port = htons(8000);

	sockaddr_in recvAddress;
	recvAddress.sin_family = AF_INET;
	recvAddress.sin_addr.s_addr = INADDR_ANY;
	recvAddress.sin_port = 0; //Not specifying a port for client so an available port will be selected when binding
	int recvAddressSize = sizeof(recvAddress);

	if (bind(clientSocketHandle, (const sockaddr*)&recvAddress, sizeof(sockaddr_in)) < 0)
	{
		std::cerr << "Failed to bind socket." << std::endl;
		return;
	}

	uint32_t packetCounter = 0;
	uint32_t recvPacketCounter = 0;
	while (true)
	{
		if (GetAsyncKeyState(VK_ESCAPE))
			return;
		uint32_t tempval = packetCounter;
		tempval = htonl(tempval);
		sendto(clientSocketHandle, (char*)(&tempval), sizeof(uint32_t), NULL, (sockaddr*)&sendtoAddress, sizeof(sockaddr_in)); //<-This works
		if (recvfrom(clientSocketHandle, (char*)(&recvPacketCounter), sizeof(uint32_t), NULL, (sockaddr*)&recvAddress, &recvAddressSize) > 0) //<-Never receiving a message
		{
			recvPacketCounter = ntohl(recvPacketCounter);
			std::cout << "Client - [Sender Port: " << ntohs(recvAddress.sin_port) << "] Received packet: " << recvPacketCounter << std::endl;
		}
		Sleep(100);
		packetCounter += 1;
	}
}

int main()
{
	char input = 0;
	while (input != 'c' && input != 's')
	{
		std::cout << "Enter 'c' for client or 's' for server:";
		std::cin >> input;
	}

	if (input == 's')
		runtype = RunType::Server;
	else
		runtype = RunType::Client;

	WSADATA init;
	if (WSAStartup(MAKEWORD(2, 2), &init) != 0)
	{
		std::cerr << "WSA Startup failed with error code: " << WSAGetLastError() << std::endl;
		return -1;
	}

	if (runtype == RunType::Server)
	{
		DoServer();
	}
	if (runtype == RunType::Client)
	{
		DoClient();
	}

	WSACleanup();

	std::cout << "Program end." << std::endl;
	return 0;
}