#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <signal.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <limits.h>
#include <time.h>

#define SEND_TIMEOUT 1
#define RECV_TIMEOUT 2
#define BUFLEN 1024 * 1024
#define MSG1 "Server very loaded.\n"
#define MSG2 "File very big.\n"
#define MSG3 "Closing connection.\n"
#define MSG4 "No such file.\n"

int iport = 8080;
int oport = 8081;
size_t max_clients = 15;
size_t max_size = 500;
size_t name_len = 15;
char *dir = ".";

char *prog_name;
int ifd;
size_t ifd_clients;

int ofd;
size_t ofd_clients;

enum {
	OUTPUT,
	INPUT
};

enum {
	WRITE,
	SEND
};

int new_server(const int port) {
	int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
	if (listen_fd < 0)
		return -1;

	struct timeval tout = {.tv_sec = RECV_TIMEOUT, .tv_usec = 0};
	if (setsockopt(listen_fd, SOL_SOCKET, SO_RCVTIMEO, &tout, sizeof(tout)) < 0)
		goto NS_ERROR;

	tout.tv_sec = SEND_TIMEOUT;
	if (setsockopt(listen_fd, SOL_SOCKET, SO_SNDTIMEO, &tout, sizeof(tout)) < 0)
		goto NS_ERROR;

	int reuse = 1;
	if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEPORT, &reuse, sizeof(reuse)) < 0)
		goto NS_ERROR;

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

	if (bind(listen_fd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0)
		goto NS_ERROR;

	if (listen(listen_fd, 0) < 0)
		goto NS_ERROR;

	return listen_fd;

NS_ERROR:
	close(listen_fd);
	return -1;
}

char *rand_string(void) {
	char abc[] = "asdfghjklzxcvbnmqwertyuiopQWERTYUIOPASDFGHJKLZXCVBNM1234567890";
	char *res = malloc(name_len + 1);
	if (res == NULL)
		return NULL;

	for (size_t i = 0; i < name_len; i++)
		res[i] = abc[rand() % (sizeof(abc) - 1)];

	res[name_len] = '\0';
	return res;
}

char *rand_name(void) {
	while (1) {
		char *filename = rand_string();
		if (filename == NULL)
			continue;

		struct stat sb;
		if (stat(filename, &sb) < 0)
			return filename;

		free(filename);
	}
}

int send_str(int fd, const char *str, ssize_t size, int flag) {
	ssize_t n = 0;
	ssize_t rbytes = size;

	while (rbytes) {
		ssize_t ret = 0;
		if (flag)
			ret = send(fd, str + n, rbytes, 0);

		else
			ret = write(fd, str + n, rbytes);

		if (ret < 0 && (errno != EAGAIN || errno != EINTR))
			return 1;

		else if (ret == 0)
			return 0;

		rbytes -= ret;
		n += size - rbytes;
	}

	return 0;
}

void upload(int cfd, pid_t par) {
	int ret = -1;
	int ffd = 0;

	char *filename = rand_name();
	if (send(cfd, filename, strlen(filename), 0) < 0)
		goto UPL_ERROR;

	if (send(cfd, "\n", 1, 0) < 0)
		goto UPL_ERROR;

	ffd = open(filename, O_CREAT | O_WRONLY, 0666);
	if (ffd < 0)
		goto UPL_ERROR;

	char buf[BUFLEN + 1];

	ssize_t tbytes = 0;
	while (1) {
		ssize_t rbytes = recv(cfd, buf, sizeof(buf), 0);
		if (rbytes < 0 && (errno != EAGAIN || errno != EINTR))
			goto UPL_ERROR;

		else if (rbytes == 0)
			break;

		tbytes += rbytes;
		if (tbytes / 1024 / 1024 >= (ssize_t)max_size) {
			send_str(cfd, MSG2, sizeof(MSG2), WRITE);
			break;
		}

		if (send_str(ffd, buf, rbytes, WRITE))
			goto UPL_ERROR;
	}

	ret = 0;

UPL_ERROR:
	if (ffd > 0)
		close(ffd);

	send_str(cfd, MSG3, sizeof(MSG3), WRITE);
	close(cfd);
	if (ret == -1)
		fprintf(stderr, "%s: input thread: %s: %s\n", prog_name, filename, strerror(errno));

	free(filename);
	kill(par, SIGUSR1);
	exit(0);
}

void load(int cfd, pid_t par) {
	int ret = -1;
	int ffd = 0;

	/* Get filename */
	char *filename = malloc(name_len + 1);
	if (filename == NULL)
		goto LD_ERROR;

	ssize_t rbytes = recv(cfd, filename, name_len, 0);
	if (rbytes <= 0) {
		if (rbytes == 0)
			ret = 0;

		goto LD_ERROR;
	}

	filename[rbytes] = '\0';
	char *ptr = strchr(filename, '\n');
	if (ptr)
		*ptr = '\0';

	if (strstr(filename, ".") || strstr(filename, "/"))
		goto LD_ERROR;

	/* INFO */
	if (!strcmp(filename, "INFO")) {
		char msg[64];
		int ret2 = snprintf(msg, sizeof(msg), "Max file size: %zuMB\nMax clients in thread: %zu\n", max_size, max_clients);

		if (send_str(cfd, msg, (ssize_t)ret2, SEND))
			goto LD_ERROR;

		ret = 0;
		goto LD_ERROR;
	}

	/* Open and send */
	ffd = open(filename, O_RDONLY);
	if (ffd < 0) {
		send_str(cfd, MSG4, sizeof(MSG4), SEND);
		goto LD_ERROR;
	}

	char buf[BUFLEN + 1];
	while (1) {
		rbytes = read(ffd, buf, sizeof(buf));
		if (rbytes < 0 && (errno != EAGAIN || errno != EINTR))
			goto LD_ERROR;

		else if (rbytes == 0)
			break;

		if (send_str(cfd, buf, rbytes, SEND))
			goto LD_ERROR;
	}

	ret = 0;

LD_ERROR:
	if (ffd > 0)
		close(ffd);

	close(cfd);
	if (ret == -1)
		fprintf(stderr, "%s: output thread: %s: %s\n", prog_name, filename, strerror(errno));

	free(filename);
	kill(par, SIGUSR2);
	exit(0);
}

void sig_handler(int sig) {
	if (sig == SIGUSR1)
		ifd_clients--;

	else if (sig == SIGUSR2)
		ofd_clients--;

	else {
		close(ifd);
		close(ofd);

		exit(0);
	}
}

void Proc(int fd, int flag) {
	static struct sockaddr_in cli_addr;
	socklen_t length = sizeof(cli_addr);

	pid_t par = getpid();

	while (1) {
		int cfd = accept(fd, (struct sockaddr *)&cli_addr, &length);
		if (cfd < 0)
			continue;

		if (flag == INPUT && ifd_clients < max_clients) {
			ifd_clients++;

			if (fork() == 0) {
				close(fd);
				upload(cfd, par);
			}
		}

		else if (flag == OUTPUT && ofd_clients < max_clients) {
			ofd_clients++;

			if (fork() == 0) {
				close(fd);
				load(cfd, par);
			}
		}

		else
			send(cfd, MSG1, sizeof(MSG1), 0);

		close(cfd);
	}
}

int main(int argc, char **argv) {
	srand(time(NULL));
	prog_name = argv[0];

	signal(SIGUSR1, sig_handler);
	signal(SIGUSR2, sig_handler);
	signal(SIGINT, sig_handler);
	signal(SIGTERM, sig_handler);
	signal(SIGCHLD, SIG_IGN);

	int opt;
	while ((opt = getopt(argc, argv, "i:o:m:d:s:n:")) != -1) {
		switch (opt) {
			case 'i':
				iport = atoi(optarg);
				break;

			case 'o':
				oport = atoi(optarg);
				break;

			case 'm':
				sscanf(optarg, "%zu", &max_clients);
				break;

			case 'd':
				dir = optarg;
				break;

			case 's':
				sscanf(optarg, "%zu", &max_size);
				break;

			case 'n':
				sscanf(optarg, "%zu", &name_len);
				break;

			default:
				printf("%s [iomds] - Simple file sharing server\n\t-i NUM Input port\tdefault: %d\n\t-o NUM Output port\tdefault: %d\n\t-m NUM Max clients\tdefault: %zu\n\t-d DIR Word directory\tdefault: %s\n\t-s NUM Max file size\tdefault: %zuMB\n\t-n NUM Filename length\tdefault: %zu\n", prog_name, iport, oport, max_clients, dir, max_size, name_len);
				return 0;
		}
	}

	if (chdir(dir) < 0) {
		fprintf(stderr, "%s: chdir: %s\n", prog_name, strerror(errno));
		return 1;
	}

	ifd = new_server(iport);
	if (ifd == -1) {
		fprintf(stderr, "%s: new_server: input socket: %s\n", prog_name, strerror(errno));
		return 1;
	}

	ofd = new_server(oport);
	if (ofd == -1) {
		fprintf(stderr, "%s: new_server: output socket: %s\n", prog_name, strerror(errno));
		return 1;
	}

	if (fork() == 0)
		Proc(ofd, OUTPUT);

	else
		Proc(ifd, INPUT);
}