#include <stdio.h>
#include <stdbool.h>
#include <windows.h>
#include <winsock.h>
#include <process.h>

#define ERRMSG_BUFFER_SIZE  100   // Size of error message buffer
#define LISTEN_PORT         1234  // Port we accept connections on
#define BUFFER_SIZE         100   // Size of buffer for sending data
#define MIN_NUMBER_TICKS    100   // Minimum number of ticks per thread
#define MAX_NUMBER_TICKS    200   // Maximum number of ticks per thread
#define SLEEP_TIME          50    // Milliseconds to sleep between ticks

// Parameters passed to new thread
struct ThreadContext
{
	SOCKET sock;  // socket used to communicate with client
	int id;       // ID of client session
	int ticks;    // number of ticks to count
};

/* Name:         NextSessionId
 * Description:  Next ID used as session ID
 * Access:       Written and read by main thread only
 */
int NextSessionId;

/* Name:         NumberOfClients
 * Description:  Current number of clients
 * Access:       Written and read by anyone
 */
int NumberOfClients;
CRITICAL_SECTION NumberOfClientsLock;

/* Name:         TotalNumberOfTicks
 * Description:  Total number of ticks spent on server
 * Access:       Written and read by anyone
 */
int TotalNumberOfTicks;
CRITICAL_SECTION TotalNumberOfTicksLock;

/* Name:         getErrorMsg
 * Desription:   Returns a human readable error message for codes from
 *               WSAGetLastError or GetLastError
 * Parameters:   code - error code to return message for
 * Return code:  Pointer to a static buffer containing the message
 * Notes:        Not thread-safe because of static buffer
 */
const char *getErrorMsg(int code)
{
	static char buffer[ERRMSG_BUFFER_SIZE];

	FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM, NULL, code, 0, buffer, sizeof(buffer), NULL);

	return buffer;
}

/* Name:         initWinsock
 * Desription:   Initialized Winsock subsystem
 * Parameters:   none
 * Return code:  true if successful
 *               false if an error occurred
 */
bool initWinsock(void)
{
	bool failed = false;

	WSADATA wsa;
	int startupResult = WSAStartup(MAKEWORD(1, 1), &wsa);
	if (startupResult != 0)
	{
		fprintf(stderr, "Couldn't start Winsock: %s\n", getErrorMsg(startupResult));
		failed = true;
	}

	return !failed;
}

/* Name:         deinitWinsock
 * Desription:   Deinitialized Winsock subsystem
 * Parameters:   none
 * Return code:  none
 */
void deinitWinsock(void)
{
	WSACleanup();
}

/* Name:         createTcpSocket
 * Description:  Creates a new, unconnected TCP socket
 * Parameters:   sock - buffer to receive socket
 * Return code:  true if successful
 *               false if an error occurred
 */
bool createTcpSocket(SOCKET *sock)
{
	bool failed = false;

	*sock = socket(AF_INET, SOCK_STREAM, 0);
	if (*sock == INVALID_SOCKET)
	{
		fprintf(stderr, "Couldn't create socket: %s\n", getErrorMsg(WSAGetLastError()));
		failed = true;
	}

	return !failed;
}

/* Name:         bindTcpSocket
 * Description:  Binds a TCP socket to a specific port on all interfaces
 * Parameters:   sock - socket to bind
 *               port - port to bind socket to
 * Return code:  true if successful
 *               false if an error occurred
 */
bool bindTcpSocket(SOCKET sock, unsigned short port)
{
	bool result = false;

	struct sockaddr_in saddr;
	saddr.sin_family = AF_INET;
	saddr.sin_port = htons(port);
	saddr.sin_addr.s_addr = INADDR_ANY;

	printf("Binding to %s:%u...", inet_ntoa(saddr.sin_addr), port);
	fflush(stdout);
	if (bind(sock, (struct sockaddr*) &saddr, sizeof(saddr)) == 0)
	{
		printf(" OK.\n");
		result = true;
	}
	else
	{
		printf(" FAILED.\n");
		fprintf(stderr, "Couldn't bind socket: %s\n", getErrorMsg(WSAGetLastError()));
	}

	return result;
}

/* Name:         setSocketToListenMode
 * Description:  Sets a bound TCP socket to listen mode
 * Parameters:   sock - socket to set to listen mode
 * Return code:  true if successful
 *               false if an error occurred
 */
bool setSocketToListenMode(SOCKET sock)
{
	bool failed = false;

	if (listen(sock, 3) == SOCKET_ERROR)
	{
		fprintf(stderr, "Couldn't set socket to listen mode: %s\n", getErrorMsg(WSAGetLastError()));
		failed = true;
	}

	return !failed;
}

/* Name:         acceptClient
 * Description:  Accepts a new client on a socket in listen mode
 * Parameters:   sock - socket in listen mode
 *               client - buffer for socket of new client
 * Return code:  true if successful
 *               false if an error occurred
 */
bool acceptClient(SOCKET sock, SOCKET *client)
{
	bool result = false;

	struct sockaddr_in saddr;
	int size = sizeof(saddr);
	*client = accept(sock, (struct sockaddr*) &saddr, &size);
	if (*client >= 0)
	{
		printf("Got a new client from %s:%u.\n", inet_ntoa(saddr.sin_addr), ntohs(saddr.sin_port));
		result = true;
	}
	else
	{
		fprintf(stderr, "Accepting client failed: %s\n", getErrorMsg(WSAGetLastError()));
	}

	return result;
}

/* Name:         safeCloseSocket
 * Description:  Checks if a socket is valid, closes it. The socket is
 *               set to an invalid value to avoid double free errors.
 * Parameters:   sock - socket to communicate with
 * Return code:  none
 */
void safeCloseSocket(SOCKET *sock)
{
	if (*sock != INVALID_SOCKET)
	{
		closesocket(*sock);
		*sock = INVALID_SOCKET;
	}
}

bool sendString(SOCKET sock, const char *str)
{
	bool failed = false;

	int len = strlen(str);
	if (send(sock, str, len, 0) != len)
	{
		fprintf(stderr, "Sending string failed: %s\n", getErrorMsg(WSAGetLastError()));
		failed = true;
	}

	return !failed;
}

void clientThread(void *arg)
{
	bool good = true;
	struct ThreadContext *ctx = arg;
	char buffer[BUFFER_SIZE];

	sprintf(buffer, "Begin client session %i\r\n", ctx->id);
	printf("%s", buffer);
	good = sendString(ctx->sock, buffer);

	for (int i = 0; good && i < ctx->ticks; i++)
	{
		EnterCriticalSection(&TotalNumberOfTicksLock);
		int ticks = ++TotalNumberOfTicks;
		LeaveCriticalSection(&TotalNumberOfTicksLock);

		EnterCriticalSection(&NumberOfClientsLock);
		int clients = NumberOfClients;
		LeaveCriticalSection(&NumberOfClientsLock);

		sprintf(buffer, "Tick %i/%i, in total %i tick%s spent, %i client%s active\r\n",
			i + 1, ctx->ticks, ticks, ticks == 1 ? "" : "s", clients, clients == 1 ? "" : "s");
		good = sendString(ctx->sock, buffer);

		Sleep(SLEEP_TIME);
	}

	sprintf(buffer, "End client session %i\r\n", ctx->id);
	printf("%s", buffer);
	if (good)
		good = sendString(ctx->sock, buffer);

	EnterCriticalSection(&NumberOfClientsLock);
	NumberOfClients--;
	LeaveCriticalSection(&NumberOfClientsLock);
	safeCloseSocket(&ctx->sock);
	free(arg);
}

int getRandomTickCount(void)
{
	int result = MIN_NUMBER_TICKS;
	int var = MAX_NUMBER_TICKS - MIN_NUMBER_TICKS + 1;
	return var ? (result + rand() % var) : result;
}

bool spawnClientThread(SOCKET sock)
{
	bool failed = false;

	struct ThreadContext *ctx = malloc(sizeof(*ctx));
	if (ctx)
	{
		ctx->sock = sock;
		ctx->id = NextSessionId++;
		ctx->ticks = getRandomTickCount();
	}
	else
	{
		fprintf(stderr, "Out of memory\n");
		failed = true;
	}

	if (!failed)
	{
		uintptr_t result = _beginthread(clientThread, 0, ctx);
		if (result == -1L)
		{
			fprintf(stderr, "Spawning thread failed: %s\n", strerror(errno));
			free(ctx);
			failed = true;
		}
	}

	return !failed;
}

/* Name:         serverMainLoop
 * Description:  Main loop of server that accepts new clients and
 *               handles them.
 * Parameters:   sock - server socket in listen mode
 * Return code:  true if successful
 *               false if an error occurred
 */
bool serverMainLoop(SOCKET sock)
{
	bool good = true;

	while (good)
	{
		SOCKET client = INVALID_SOCKET;
		good = acceptClient(sock, &client);

		EnterCriticalSection(&NumberOfClientsLock);
		NumberOfClients++;
		LeaveCriticalSection(&NumberOfClientsLock);

		if (good)
			good = spawnClientThread(client);

		if (!good)
		{
			EnterCriticalSection(&NumberOfClientsLock);
			NumberOfClients--;
			LeaveCriticalSection(&NumberOfClientsLock);
			safeCloseSocket(&client);
		}
	}

	return good;
}

int main(void)
{
	InitializeCriticalSection(&NumberOfClientsLock);
	InitializeCriticalSection(&TotalNumberOfTicksLock);

	bool startupSuccessful = false;
	bool good = startupSuccessful = initWinsock();

	SOCKET sock = INVALID_SOCKET;
	if (good)
		good = createTcpSocket(&sock);

	if (good)
		good = bindTcpSocket(sock, LISTEN_PORT);

	if (good)
		good = setSocketToListenMode(sock);

	if (good)
		good = serverMainLoop(sock);

	safeCloseSocket(&sock);

	if (startupSuccessful)
		deinitWinsock();

	DeleteCriticalSection(&NumberOfClientsLock);
	DeleteCriticalSection(&TotalNumberOfTicksLock);

	return good ? EXIT_SUCCESS : EXIT_FAILURE;
}