#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 MSG1 "Server is very busy.\n"
#define MSG2 "File is very large.\n"
#define MSG3 "No such file.\n"
#define MSG4 "No namespace available.\n"

#ifndef VERSION
#define VERSION "1.2.3"
#endif

typedef struct {
	/* Clients on thread */
	size_t max_clients;

	/* Max file size */
	size_t max_size;

	/* Buffer size for thread */
	size_t buf_size;

	/* Max filename lenght */
	size_t name_len;

	/* Directory for files */
	char *dir;

	/* Message of the day */
	char *motd;
} CFG;
CFG cfg = {.max_clients = 15, .max_size = 1, .buf_size = 32000, .name_len = 15, .dir = "."};

typedef struct {
	int fd;
	char *buf;
	char *filename;
} CLIENT_THREAD;

int ifd;
int ofd;
size_t clients;

enum {
	OUTPUT,
	INPUT
};

enum {
	WRITE,
	SEND
};

int new_server(const int port, const int backlog) {
	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 NEW_SOCKET_ERROR;

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

	int reuse = 1;
	if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEPORT, &reuse, sizeof(reuse)) < 0)
		goto NEW_SOCKET_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 NEW_SOCKET_ERROR;

	if (listen(listen_fd, backlog) < 0)
		goto NEW_SOCKET_ERROR;

	return listen_fd;

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

char rand_name(char *buf) {
	const char abc[] = "asdfghjklzxcvbnmqwertyuiopQWERTYUIOPASDFGHJKLZXCVBNM1234567890";

	srand(time(NULL));
	for (int i = 0; i < 1000; i++) {
		for (size_t i = 0; i < cfg.name_len; i++)
			buf[i] = abc[rand() % (sizeof(abc) - 1)];

		buf[cfg.name_len] = '\0';

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

	return 1;
}

int send_str(int fd, char *str, ssize_t size, int flag) {
	while (size) {
		ssize_t ret = 0;
		if (flag)
			ret = send(fd, str, size, MSG_NOSIGNAL);

		else
			ret = write(fd, str, size);

		if (ret < 0)
			return 1;

		else if (ret == 0)
			return 0;

		size -= ret;
		str += ret;
	}

	return 0;
}

void upload(CLIENT_THREAD *CT) {
	if (rand_name(CT->filename)) {
		send(CT->fd, MSG4, sizeof(MSG4), 0);
		return;
	}

	send(CT->fd, CT->filename, strlen(CT->filename), 0);
	if (send(CT->fd, "\n", 1, 0) < 0)
		return;

	/* Create file */
	int ffd = open(CT->filename, O_CREAT | O_WRONLY, 0666);
	if (ffd < 0)
		goto UPLOAD_ERROR;

	size_t tbytes = 0;
	while (1) {
		ssize_t rbytes = recv(CT->fd, CT->buf, cfg.buf_size, MSG_NOSIGNAL);
		if (rbytes < 0 && errno != EINTR)
			goto UPLOAD_ERROR;

		else if (rbytes == 0)
			break;

		tbytes += (size_t)rbytes;
		if (tbytes / 1024 / 1024 >= cfg.max_size) {
			send_str(CT->fd, MSG2, sizeof(MSG2), WRITE);
			break;
		}

		if (send_str(ffd, CT->buf, rbytes, WRITE))
			goto UPLOAD_ERROR;
	}

UPLOAD_ERROR:
	if (ffd != -1)
		close(ffd);
}

void load(CLIENT_THREAD *CT) {
	ssize_t rbytes = recv(CT->fd, CT->filename, cfg.name_len, 0);
	if (rbytes <= 0)
		return;

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

	if (strstr(CT->filename, ".") || strstr(CT->filename, "/"))
		return;

	/* INFO */
	if (!strcmp(CT->filename, "INFO")) {
		char msg[336];
		int ret = snprintf(msg, sizeof(msg), "Max file size:         %zuMB\nMax clients:      %zu\nBuffer size:           %zu\nMOTD: %s\n", cfg.max_size, cfg.max_clients, cfg.buf_size, (cfg.motd) ? cfg.motd : "");

		send_str(CT->fd, msg, (ssize_t)ret, SEND);
		return;
	}

	/* Open */
	int ffd = open(CT->filename, O_RDONLY);
	if (ffd < 0) {
		send_str(CT->fd, MSG3, sizeof(MSG3), SEND);
		goto LOAD_ERROR;
	}

	while (1) {
		rbytes = read(ffd, CT->buf, cfg.buf_size);
		if (rbytes < 0 && errno != EINTR)
			goto LOAD_ERROR;

		else if (rbytes == 0)
			break;

		if (send_str(CT->fd, CT->buf, rbytes, SEND))
			goto LOAD_ERROR;
	}

LOAD_ERROR:
	if (ffd != -1)
		close(ffd);
}

void sig_handler(int sig) {
	if (sig == SIGUSR1 && clients)
		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 (clients < cfg.max_clients) {
			clients++;

			if (fork() == 0) {
				close(fd);

				/* Allocate buffers */
				char *buf = malloc(cfg.buf_size);
				char *filename = malloc(cfg.name_len + 1);
				if (filename == NULL || buf == NULL) {
					if (filename != NULL)
						free(filename);

					if (buf != NULL)
						free(buf);

					exit(1);
				}

				/* Start functions */
				CLIENT_THREAD CT = {.fd = cfd, .buf = buf, .filename = filename};
				if (flag == INPUT)
					upload(&CT);

				else if (flag == OUTPUT)
					load(&CT);

				/* Clean */
				sleep(1);

				free(buf);
				close(cfd);
				kill(par, SIGUSR1);
				exit(0);
			}
		}

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

		close(cfd);
	}
}

int main(int argc, char **argv) {
	signal(SIGUSR1, sig_handler);
	signal(SIGINT, sig_handler);
	signal(SIGTERM, sig_handler);
	signal(SIGCHLD, SIG_IGN);
	signal(SIGPIPE, SIG_IGN);

	int backlog = 1;
	int iport = 8080;
	int oport = 8081;
	char daemon_flag = 0;

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

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

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

			case 'd':
				cfg.dir = optarg;
				break;

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

			case 'n':
				sscanf(optarg, "%zu", &cfg.name_len);
				if (cfg.name_len < 4) {
					fprintf(stderr, "%s: name lenght can not be less than 4\n", argv[0]);
					return 1;
				}

				break;

			case 'b':
				sscanf(optarg, "%zu", &cfg.buf_size);
				break;

			case 't':
				if (strlen(optarg) > 256) {
					fprintf(stderr, "%s: MOTD very long\n", argv[0]);
					return 1;
				}

				cfg.motd = optarg;
				break;

			case 'l':
				backlog = atoi(optarg);
				break;

			case 'D':
				daemon_flag = 1;
				break;

			case 'V':
				puts("Version: " VERSION "\nWritten under WTFPL License.");
				return 0;

			default:
				printf("%s [iomdsnbtlDV] - 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 STR Working dir\tdefault: %s\n\t-s NUM Max file size\tdefault: %zuMB\n\t-n NUM Filename length\tdefault: %zu\n\t-b NUM Buffer size\tdefault: %zuB\n\t-t STR Set MOTD\n\t-l listen() backlog\tdefault: %d\n\t-D Run as daemon\n\t-V Print version and license\n", argv[0], iport, oport, cfg.max_clients, cfg.dir, cfg.max_size, cfg.name_len, cfg.buf_size, backlog);
				return 0;
		}
	}

	if (daemon_flag)
		if (daemon(1, 1) < 0) {
			fprintf(stderr, "%s: daemon: %s\n", argv[0], strerror(errno));
			return 1;
		}

	if (chdir(cfg.dir) < 0) {
		fprintf(stderr, "%s: chdir: %s\n", argv[0], strerror(errno));
		return 1;
	}

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

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

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

	else
		Proc(ifd, INPUT);

	return 0;
}