#include <stdio.h>
#include <netdb.h>
#include <netinet/in.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>
#include <ctype.h>

#include <pthread.h>
#include <stdbool.h>
#include <stdatomic.h>

#define MAX_CLIENTS 50

#define MESG_SIZE 80
#define SA struct sockaddr

const int Distance = 'a'-'A';

int firstClientPort;

typedef _Atomic(uint64_t) atomic_uint64_t;

// Bitset of used client slots (and thus ports)
// 1 = Taken, 0 = Free
atomic_uint64_t freeCSlotBits = 0;

// Count trailing zeros
// if x == 0 returns 64
uint32_t safe_ctz(uint64_t x) {
	return x == 0 ? 64 : __builtin_ctz(x);
}

bool reserveFreeCSlot(uint32_t* slot) {
	for (;;) {
		const uint64_t prev = freeCSlotBits;
		uint32_t freeSlot = safe_ctz(~prev);
		
		if (freeSlot < MAX_CLIENTS) {
			uint64_t new = prev | (0x1 << freeSlot);
			
			uint64_t expected = prev;
			if (atomic_compare_exchange_strong(&freeCSlotBits, &expected, new)) {
				*slot = freeSlot;
				return true;
			} else {
				// Try again
				continue;
			}
		} else {
			return false;
		}
	}
}

void freeCSlot(uint32_t slot) {
	if (slot < MAX_CLIENTS) {
		atomic_fetch_and(&freeCSlotBits, ~(0x1 << slot));
	}
}

typedef struct {
	int connfd;
	int slot;
} ClientData;

// Function designed for chat between client and server.
void changeCase(char* _str) {
	while (*_str != 0) {
		if(*_str >= 'a'  && *_str <= 'z') {
			*_str -= Distance;
		} else if(*_str >= 'A'  && *_str <= 'Z') {
			*_str += Distance;
		}
		_str++;
	}
}

void serverFunction(int sockfd) {
	char buffer[MESG_SIZE];

	// infinite loop for chat
	while (1) {
		// clear buffer
		bzero(buffer, MESG_SIZE);
		// read the message from client and copy it in buffer
		if (read(sockfd, buffer, sizeof(buffer)) <= 0) {
			printf("Failed to read client conn...\n");
			break;
		}
		
		// exchange upper-case letters by lower-case letter and vice versa.
		changeCase(buffer);
		
		// and send that buffer to client
		write(sockfd, buffer, sizeof(buffer));
		
		// if msg contains "QUIT" then server exit and chat ended.
		if (strncmp("QUIT", buffer, 4) == 0) {
			printf("Server Exit...\n");
			break;
		}
	}
}

/*
int setupConnection(int _port) {
	int sockfd, connfd, len;
	struct sockaddr_in servaddr, client;
	
	// socket create and verification
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd == -1) {
		fprintf(stderr, "Error: Cannot create socket --- exit\n");
		exit(3);
	} else {
		printf("Socket created.\n");
	}
	bzero(&servaddr, sizeof(servaddr));
	
	// set up socket
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
	servaddr.sin_port = htons(_port);
	
	if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &(int){1}, sizeof(int)) == -1) {
		fprintf(stderr, "Error: Failed to set SO_REUSEADDR --- exit\n");
		exit(14);
	}
    
	// Binding socket to any IP
	if ((bind(sockfd, (SA*)&servaddr, sizeof(servaddr))) != 0) {
		fprintf(stderr, "Error: Cannot bind socket --- exit\n");
		exit(4);
	} else {
		printf("Socket bound.\n");
	}
	
	// Server listens
	if ((listen(sockfd, 5)) != 0) {
		fprintf(stderr, "Error: Cannot listen --- exit\n");
		exit(5);
	} else { 
		printf("Server listening.\n");
	}
	len = sizeof(client);

	// Accept data packet from client
	connfd = accept(sockfd, (SA*)&client, &len);
	if (connfd < 0) {
		fprintf(stderr, "Error: Server accept failed --- exit\n");
		exit(6);
	} else {
		printf("Server accept client.\n");
	}

	return connfd;
}
*/

int setupClientSocket(int* sockPort) {
	int sockfd, len;
	struct sockaddr_in servaddr, client;
	
	// socket create and verification
	sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd == -1) {
		fprintf(stderr, "Error: Cannot create socket --- exit\n");
		exit(3);
	} else {
		printf("Socket created.\n");
	}
	bzero(&servaddr, sizeof(servaddr));
	
	// set up socket
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
	servaddr.sin_port = htons(0); // Bind to any free port
	
	// Binding socket to any IP
	if ((bind(sockfd, (SA*)&servaddr, sizeof(servaddr))) != 0) {
		fprintf(stderr, "Error: Cannot bind socket (setupClientConnection) --- exit\n");
		exit(4);
	} else {
		printf("Socket bound.\n");
	}
	
	// Server listens
	if ((listen(sockfd, 5)) != 0) {
		fprintf(stderr, "Error: Cannot listen --- exit\n");
		exit(5);
	} else { 
		printf("Server listening.\n");
	}
	
	// Get actual port of new socket
	struct sockaddr_in sin;
	int sinLen = sizeof(sin);
	getsockname(sockfd, (SA*)&sin, &sinLen);
	*sockPort = ntohs(sin.sin_port);
	
	return sockfd;
}

int finalizeClientConnection(int sockfd) {
	// Accept data packet from client
	SA client;
	int len = sizeof(SA);
	int connfd = accept(sockfd, (SA*)&client, &len);
	if (connfd < 0) {
		fprintf(stderr, "Error: Server accept failed --- exit\n");
		exit(6);
	} else {
		printf("Server accept client.\n");
	}
	return connfd;
}

int setupListenSocket(int _port) {
	struct sockaddr_in servaddr;

	// socket create and verification
	int sockfd = socket(AF_INET, SOCK_STREAM, 0);
	if (sockfd == -1) {
		fprintf(stderr, "Error: Cannot create socket --- exit\n");
		exit(3);
	} else {
		printf("Socket created.\n");
	}
	bzero(&servaddr, sizeof(servaddr));

	// set up socket
	servaddr.sin_family = AF_INET;
	servaddr.sin_addr.s_addr = htonl(INADDR_ANY);
	servaddr.sin_port = htons(_port);

	// Binding socket to any IP
	if ((bind(sockfd, (SA*)&servaddr, sizeof(servaddr))) != 0) {
		fprintf(stderr, "Error: Cannot bind socket --- exit\n");
		exit(4);
	} else {
		printf("Socket bound.\n");
	}
	
	// Server listens
	if ((listen(sockfd, 5)) != 0) {
		fprintf(stderr, "Error: Cannot listen --- exit\n");
		exit(5);
	} else { 
		printf("Server listening.\n");
	}
	
	return sockfd;
}

void firstTouch(int sockfd, int connPort) {
	char buffer[MESG_SIZE];
	bzero(buffer, MESG_SIZE);
	read(sockfd, buffer, MESG_SIZE);
	printf("received: %s\n", buffer);
	((int*)buffer)[0] = connPort;
	write(sockfd, buffer, MESG_SIZE);
	close(sockfd);
}

void permantConnection(int slot, int connfd) {
	serverFunction(connfd);
	close(connfd);
	
	printf("Client #%d disconnected.\n", slot);
	
	// Clean up, free client slot
	freeCSlot(slot);
}

void* handleClientConnection(void* args) {
	ClientData* client = (ClientData*)args;
	
	permantConnection(client->slot, client->connfd);
	free(client);
}

bool tryAcceptClient(int listenSockFd) {
	// Prereserve client slot
	uint32_t clientSlot;
	if (!reserveFreeCSlot(&clientSlot)) {
		// No client slots free
		return false;
	}
	
	// Accept new client
	int clientAddrLen;
	struct sockaddr clientAddr;
	
	int connfd = accept(listenSockFd, (SA*)&clientAddr, &clientAddrLen);
	if (connfd < 0) {
		fprintf(stderr, "Error: Server accept failed --- exit\n");
		exit(6);
	} else {
		printf("Server accept client #%d.\n", clientSlot);
	}
	
	// Setup actual conn
	int laterPort;
	int laterSockFd = setupClientSocket(&laterPort);
	
	firstTouch(connfd, laterPort);
	
	int clientConnFd = finalizeClientConnection(laterSockFd);
	close(laterSockFd);
	
	ClientData* client = malloc(sizeof(ClientData));
	client->connfd = clientConnFd;
	client->slot = clientSlot;
	
	pthread_t clientThread;
	pthread_create(&clientThread, NULL, handleClientConnection, (void*)client);
	
	return true;
}

int main(int argc, char** argv) {
	int listen_port = 4567;
	
	if (argc > 2) {
		fprintf(stderr, "usage: %s [port] --- exit\n", argv[0]);
		return 1;
	}
	if(argc == 2) {
		char* tmp = argv[1];
		while(tmp[0] != 0) {
			if(!isdigit((int)tmp[0])) {
				fprintf(stderr, "Error: %s is no valid port --- exit\n", argv[1]);
				return 2;
			}
			tmp++;
		}
		listen_port = atoi(argv[1]);
	}
	
	// Set first client port as listenport + 1
	firstClientPort = listen_port + 1;
	
	// Setup server listen socket
	int listenSockFd = setupListenSocket(listen_port);
	
	for(;;) {
		if (!tryAcceptClient(listenSockFd)) {
			// No slots free, wait 1ms to try again
			usleep(1000);
		}
	}
	return 0;
}