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

#include<vector>
using namespace std;

char buffer[1024];

struct client_info {
	int fd;
	struct sockaddr_in address;
	char name[32];
};

// Don't forget to free the returned pointer!
char *get_message(int fd){
	uint16_t mlen;
	if(read(fd, &mlen, 2) != 2)
		return 0;
	// TODO:  Explain need for cast here
	char* message = (char*)malloc(mlen+1);
	if(recv(fd, message, mlen, MSG_WAITALL) != mlen){
		free(message);
		return 0;
	}
	message[mlen] = 0;
	return message;
}

// Don't forget to free the returned pointer!
void* make_nbm(char *sender, char *message, uint16_t *len){
	uint16_t mlen = strlen(message);
	*len = 35 + mlen;
	char *nbm = (char*)malloc(*len);
	nbm[0] = 9;
	strncpy(nbm + 1, sender, 32);
	memcpy(nbm + 33, &mlen, 2);
	strncpy(nbm+35, message, mlen);
	return nbm;
}

// Returns 1 on failure, 0 on success.  Same convention as main.
int send_error(int fd, const char *message){
	uint16_t len;
	len = 2;
	if(1 != write(fd, &len, 1)) // Just the first byte of the uint16_t
		return 1;
	len = strlen(message);
	if(2 != write(fd, &len, 2)) // 2 bytes (the whole thing) this time
		return 1;
	if(len != write(fd, message, len))
		return 1;
	return 0;
}

vector<struct client_info*> clients;

void* handle_client(void* param){
	// TODO:  Explain cast here too, although it is the same problem as above
	struct client_info *client = (struct client_info*)param;
	clients.push_back(client);

	printf("Connection made from address %s with fd %d\n", inet_ntoa(client->address.sin_addr), client->fd);
	
	// Log in procedure
	char type;
	while(1){
		if(!read(client->fd, &type, 1))
			goto disconnect;
		if(type != 1){
			send_error(client->fd, "Expected login message");
			continue;
		}
		recv(client->fd, client->name, 32, MSG_WAITALL);
		printf("New Client:  %s\n", client->name);
		type = 4; // TODO:  Check and see if that username was already used
		write(client->fd, &type, 1);
		break;
	}
	
	// This section assumes the login worked
	type = 3;
	for(auto c : clients){
		write(c->fd, &type, 1);
		write(c->fd, client->name, 32);
		printf("Sent announcement to client %s for join of %s\n", c->name, client->name);
	}

	// Start main loop looking for messages
	size_t readlen;
	while(1){
		char next_message_type;
		if(!read(client->fd, &next_message_type, 1))
			break;
		switch(next_message_type){
			case 1:
				// Change username
				// Leave this one for later
				// Only the server sends 2, 3, 4, 5, and 8
			case 6:{
				// Broadcast message
				uint16_t len;
				char* message = get_message(client->fd);
				void* nbm = make_nbm(client->name, message, &len);
				printf("Relaying broadcast message %s from %s\n", message, client->name);
				for(struct client_info* c : clients)
					write(c->fd, nbm, len);
				free(message);
				free(nbm);
				break;
			}
			case 7:{
				// Specific Message
				uint16_t len;
				char* message = get_message(client->fd);
				void* nsm = make_nbm(client->name, message, &len);
				((char*)nsm)[0] = 7;
				for(struct client_info* c : clients)
					if(!strcmp(c->name, client->name)){
						write(c->fd, nsm, len);
						goto success;
					}
				send_error(client->fd, "No such user!");
success:		
				free(message);
				free(nsm);
				break;
			}
		}
	}
	for(int i = 0; i < clients.size(); i++){
		if(clients[i] == client){
			clients[i] = clients[clients.size() - 1];
			clients.pop_back();
		}
	}
	char disconnect_message[33];
	disconnect_message[0] = 9;
	strcpy(disconnect_message+1, client->name);
	for(struct client_info* c : clients)
		write(c->fd, disconnect_message, 33);
disconnect:
	printf("Done with connection from address %s, client %s\n", inet_ntoa(client->address.sin_addr), client->name);
	close(client->fd);
	free(client);
	return 0;
}

int main(int argc, char ** argv){
	struct sockaddr_in sad;
	sad.sin_port = htons(5141);
	sad.sin_addr.s_addr = INADDR_ANY;
	sad.sin_family = AF_INET;

	int skt = socket(AF_INET, SOCK_STREAM, 0);
	bind(skt, (struct sockaddr *)(&sad), sizeof(struct sockaddr_in));
	listen(skt, 5);
	socklen_t address_size = sizeof(struct sockaddr_in);
	for(;;){
		struct client_info *new_client = (struct client_info*)malloc(sizeof(struct client_info));
		new_client->fd = accept(skt, (struct sockaddr *)(&(new_client->address)), &address_size);
		// Start a new thread to deal with the new client
		pthread_t T;
		pthread_create(&T, 0, handle_client, new_client);
	}
	return 0;
}