#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/select.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>

#define LISTEN_PORT  1234  // Port we accept connections on
#define BUFFER_SIZE  1024  // Size of buffer for receiving data
#define MAX_CLIENTS  3     // Maximum number of clients

/* 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(int *sock)
{
	bool failed = false;

	*sock = socket(AF_INET, SOCK_STREAM, 0);
	if (*sock == -1)
	{
		fprintf(stderr, "Couldn't create socket: %s\n", strerror(errno));
		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(int 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", strerror(errno));
	}

	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(int sock)
{
	bool failed = false;

	if (listen(sock, 3) == -1)
	{
		fprintf(stderr, "Couldn't set socket to listen mode: %s\n", strerror(errno));
		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(int sock, int *client)
{
	bool result = false;

	struct sockaddr_in saddr;
	socklen_t 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", strerror(errno));
	}

	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(int *sock)
{
	if (*sock != -1)
	{
		close(*sock);
		*sock = -1;
	}
}

/* Name:         initClientList
 * Description:  Initializes client list
 * Parameters:   list - array of sockets to initialize
 *               size - number of elements in array
 * Return code:  none
 */
void initClientList(int *list, int size)
{
	for (int i = 0; i < size; i++)
		list[i] = -1;
}

/* Name:         addClientToList
 * Description:  Adds a client to the client list
 * Parameters:   list - array of sockets
 *               size - number of elements in array
 *               client - socket to add
 *               slot - buffer to receive slow of new client
 * Return code:  true if successful
 *               false if list is full
 */
bool addClientToList(int *list, int size, int client, int *slot)
{
	bool added = false;

	for (int i = 0; !added && i < size; i++)
	{
		if (list[i] == -1)
		{
			list[i] = client;
			*slot = i;
			added = true;
		}
	}

	return added;
}

/* Name:         cleanupClientList
 * Description:  Closes all clients in list
 * Parameters:   list - array of sockets to close
 *               size - number of elements in array
 * Return code:  none
 */
void cleanupClientList(int *list, int size)
{
	for (int i = 0; i < size; i++)
		safeCloseSocket(&list[i]);
}

/* Name:         addSocketsToSet
 * Description:  Adds a list of sockets to a fd_set
 * Parameters:   list - array of sockets to add to set
 *               size - number of elements in array
 *               set - set to add sockets to
 *               maxfd - highest numbered socket in set
 * Return code:  none
 */
void addSocketsToSet(int *list, int size, fd_set *set, int *maxfd)
{
	for (int i = 0; i < size; i++)
	{
		if (list[i] != -1)
		{
			FD_SET(list[i], set);
			if (list[i] > *maxfd)
				*maxfd = list[i];
		}
	}
}

/* Name:         handleServerReadable
 * Description:  Handles a readable server socket by accepting a new
 *               client which is added to the client list.
 * Parameters:   server - server socket to handle
 *               clients - array of client sockets
 *               size - number of elements in array
 * Return code:  true if successful
 *               false if an error occurred
 */
bool handleServerReadable(int server, int *clients, int size)
{
	printf("----- Server socket -----\n");

	int newClient;
	bool good = acceptClient(server, &newClient);

	if (good)
	{
		int slot;
		if (addClientToList(clients, size, newClient, &slot))
		{
			printf("Client added to list into slot %i.\n", slot);
		}
		else
		{
			safeCloseSocket(&newClient);
			printf("Out of client slots, new client dropped.\n");
		}
	}

	return good;
}

/* Name:         handleClientReadable
 * Description:  Handles readable client socket(s) by receiving data
 *               or removing clients if connection was closed.
 * Parameters:   clients - array of client sockets
 *               size - number of elements in array
 *               set - set of readable client sockets
 * Return code:  true if successful
 *               false if an error occurred
 */
bool handleClientReadable(int *clients, int size, fd_set *set)
{
	bool good = true;

	for (int i = 0; good && i < size; i++)
	{
		if (FD_ISSET(clients[i], set))
		{
			char buffer[BUFFER_SIZE];
			printf("----- Client slot %i -----\n", i);
			int bytes = recv(clients[i], buffer, sizeof(buffer) - 1, 0);
			if (bytes > 0)
			{
				buffer[bytes] = '\0';
				printf("Received %i bytes: <%s>\n", bytes, buffer);
			}
			else if (bytes == 0)
			{
				printf("Client closed connection.\n");
				safeCloseSocket(&clients[i]);
			}
			else if (bytes == -1)
			{
				fprintf(stderr, "Receiving data from client failed: %s\n", strerror(errno));
				safeCloseSocket(&clients[i]);
				good = false;
			}
		}
	}

	return good;
}

/* 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(int sock)
{
	bool good = true;
	int clients[MAX_CLIENTS];
	initClientList(clients, MAX_CLIENTS);

	while (good)
	{
		fd_set readfds;
		FD_ZERO(&readfds);

		int maxfd = 0;
		addSocketsToSet(clients, MAX_CLIENTS, &readfds, &maxfd);
		addSocketsToSet(&sock, 1, &readfds, &maxfd);

		int count = select(maxfd + 1, &readfds, NULL, NULL, NULL);
		if (count > 0)
		{
			if (FD_ISSET(sock, &readfds))
				good = handleServerReadable(sock, clients, MAX_CLIENTS);
			else
				good = handleClientReadable(clients, MAX_CLIENTS, &readfds);
		}
		else
		{
			fprintf(stderr, "Waiting for events failed: %s\n", strerror(errno));
			good = false;
		}
	}

	cleanupClientList(clients, MAX_CLIENTS);
	return good;
}

int main(void)
{
	int sock = -1;
	bool good = createTcpSocket(&sock);

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

	if (good)
		good = setSocketToListenMode(sock);

	if (good)
		good = serverMainLoop(sock);

	safeCloseSocket(&sock);

	return good ? EXIT_SUCCESS : EXIT_FAILURE;
}