#include<sys/socket.h>
#include<sys/types.h>
#include<netinet/ip.h>
#include<arpa/inet.h>
#include<stdio.h>
#include<unistd.h>
#include<signal.h>
#include<stdlib.h>
#include<string.h>
#include<pthread.h>
#include<vector>

/*
 * Name:
 * 	type (1 byte):  1
 * 	name length (2 bytes, unsigned 16-bit):  length of the name
 * 	new_name (name length bytes, some kind of text):  the actual name
 *
 * Message:
 * 	type (1 byte):  2
 * 	message length (4 bytes, uint32):  length of the message
 * 	actual message (message length bytes):  actual message
 */

int skt; // Global.  This is the socket that the server will listen on
#define MAX_NAMELEN 128

std::vector<int> all_clients;

void* handle_client(void* arg){
	int client_fd = (int)(size_t)arg;
	char name[128] = "default";
	ssize_t length;
	uint16_t name_length = 7;
	while(1){
		char type;
		length = read(client_fd, &type, 1);
		if(length != 1)
			break;
		if(type == 1){ // change name
			length = recv(client_fd, &name_length, 2, MSG_WAITALL);
			if(length != 2)
				break;
			if(name_length > MAX_NAMELEN){ // We could just have read the first part of the name instead
				break;
			}
			length = recv(client_fd, name, name_length, MSG_WAITALL);
			name[length] = 0;
			if(length != name_length)
				break;
			printf("Name change, now %s\n", name);			
		} else if (type == 2){ // send message
			uint32_t message_length;
			length = recv(client_fd, &message_length, 4, MSG_WAITALL);
			if(length != 4)
				break;
			char *message_buffer = (char*)malloc(message_length + 2 + name_length + 4);
			length = recv(client_fd, message_buffer + 2 + name_length + 4, message_length, MSG_WAITALL);
			memcpy(message_buffer, &name_length, 2);
			memcpy(message_buffer + 2, name, name_length);
			memcpy(message_buffer + 2 + name_length, &message_length, 4);
			// relay the message to all the clients
			for(auto cfd : all_clients)
				if(cfd != client_fd)
					write(cfd, message_buffer, message_length + 2 + name_length + 4);
			free(message_buffer);		
		} else {
			printf("Unknown type:  %d\n", type);
		}	
	}
	for(int i = 0; i < all_clients.size(); i++)
		if(all_clients[i] == client_fd){ // We found our own entry
			all_clients[i] = all_clients[all_clients.size() - 1];
			all_clients.pop_back();
		}
	printf("Closing connection for %s\n", name);
	close(client_fd);
	return 0;
}

void close_server(int signal_number){
	printf("Closing down the server\n");
	close(skt);
}

/*
 * Usage:
 * 	server [-d] [port_number]
 */

int main(int argc, char ** argv){

	int port = 5142;
	char daemon = 0;
	for(int i = 1; i < argc; i++){
		if(0 == strcmp(argv[i], "-d"))
			daemon = 1;
		else if(0 == strcmp(argv[i], "--help")){
			printf("Usage:  %s [-d] [port]\n");
			return 0;
		}
		else
			port = atoi(argv[i]);
	}
	if(!port){
		printf("Invalid Port\n");
		return 1;
	}
	printf("Starting server on port %d\n", port);

	if(daemon){
		pid_t pid = fork();
		if(pid) // We're the parent process
			return 0;
	}

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

	skt = socket(AF_INET, SOCK_STREAM, 0);
	if(skt == -1){
		perror("socket");
		return 1;
	}
	struct sigaction close_action;
	close_action.sa_handler = close_server;
	if(sigaction(SIGINT, &close_action, 0)){
		perror("sigaction");
		return 1;
	}

	if( bind(skt, (struct sockaddr *)(&sad), sizeof(struct sockaddr_in)) ){
		perror("bind");
		return 1;
	}
	if( listen(skt, 5) ){
		perror("listen");
		return 1;
	}	
	int client_fd;
	struct sockaddr_in client_address;
	socklen_t address_size = sizeof(struct sockaddr_in);
	for(;;){
		client_fd = accept(skt, (struct sockaddr *)(&client_address), &address_size);
		if(client_fd == -1){
			perror("accept");
			break;
		}
		// new thread for the new client
		pthread_t t;
		pthread_create(&t, 0, handle_client, (void*)(size_t)client_fd);
		all_clients.push_back(client_fd);
		printf("Connection made from address %s\n", inet_ntoa(client_address.sin_addr));
	}
	return 0;
}