#ifdef __linux__
#define _POSIX_C_SOURCE 200809L
#define _DEFAULT_SOURCE
#endif

#include <stdlib.h>
#include <stdio.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <string.h>
#include <syslog.h>
#include <limits.h>
#include <errno.h>
#include <unistd.h>
#ifdef __FreeBSD__
#include <sys/param.h>
#endif

#include "handlers.h"
#include "packet.h"
#include "netascii.h"

inline nbd_opmode get_mode(char *mode) {
	nbd_opmode opmode;
	if (strncmp(mode, (char *)&"netascii", NBD_NBTPD_ARGS_MODE_MAX) == 0) {
		opmode = NETASCII;
	} else if (strncmp(mode, (char *)&"octet", NBD_NBTPD_ARGS_MODE_MAX) == 0) {
		opmode = OCTET;
	} else if (strncmp(mode, (char *)&"mail", NBD_NBTPD_ARGS_MODE_MAX) == 0) {
		opmode = MAIL;
	} else {
		opmode = NOT_FOUND;
	}
	return opmode;
}

int makesock(nbd_nbtpd_args *argptr) {
	int s = socket(AF_INET, SOCK_DGRAM, 0);
	if (s <= 0) {
		syslog(LOG_ERR, "unable to define socket: %s", strerror(errno));
		return -1;
	}
	struct timeval timeout = { 30, 0 };
	if (setsockopt(s, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout)) < 0) {
		syslog(LOG_ERR, "unable to set socket timeout: %s", strerror(errno));
		close(s);
		return -1;
	}
	if (connect(s, (struct sockaddr *)&(argptr->client), sizeof(struct sockaddr)) < 0) {
		syslog(
			LOG_ERR,
			"unable to connect to client %s: %s",
			inet_ntoa(argptr->client.sin_addr),
			strerror(errno)
		);
		close(s);
		return -1;
	}
	return s;
}

char *checkpath(nbd_nbtpd_args *argptr, uint8_t read) {
	char *fname = NULL, *wd = NULL, *path = NULL, *ptr = NULL;
	char **parts = NULL;
	if (!is_netascii_str((char *)&argptr->path)) {
		argptr->err = 1;
		goto cleanup;
	}
	if (read) {
		fname = realpath((char *)&argptr->path, NULL);
	} else {
		//TODO: figure out how to canonicalize non-existent path
		char *tok = NULL;
		parts = malloc(sizeof(parts) * NBD_NBTPD_ARGS_PATH_MAX);
		if (parts == NULL) {
			goto cleanup;
		}
		memset(parts, '\0', sizeof(parts) * NBD_NBTPD_ARGS_PATH_MAX);
		ptr = (char *)&argptr->path;
		for (int i = 0; ((tok = strsep(&ptr, "/")) != NULL); i++ ) {
			parts[i] = tok;
		}
		path = malloc(NBD_NBTPD_ARGS_PATH_MAX);
		if (path == NULL) {
			goto cleanup;
		}
		memset(path, '\0', NBD_NBTPD_ARGS_PATH_MAX);
		int z = 0;
		for (int i = 0; parts[i] != NULL; i++) {
			if (strncmp(parts[i], "..", 2) == 0) {
				continue;
			}
			if (strncmp(parts[i], ".", NBD_NBTPD_ARGS_PATH_MAX) == 0) {
				continue;
			}
			for (int x = 0; parts[i][x] != '\0'; x++) {
				if (z < NBD_NBTPD_ARGS_PATH_MAX) {
					path[z++] = parts[i][x];
				} else {
					goto cleanup;
				}
			}
			if (z < NBD_NBTPD_ARGS_PATH_MAX) {
				path[z++] = '/';
			} else {
				goto cleanup;
			}
		}
		// erase trailing slash
		path[z - 1] = '\0';
		fname = strdup(path);
	}
	if (fname == NULL) {
		syslog(LOG_ERR, "unable to get real path: %s", strerror(errno));
		argptr->err = 1;
		goto cleanup;
	}
/* * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
 * POSIX does not define what happens when getcwd() is given NULL.     *
 * This means that illumos libc and GNU libc both return a newly       *
 * malloc()'d string, while FreeBSD libc does not.  Also, FreeBSD uses *
 * constant MAXPATHLEN defined in <sys/param.h> rather than PATH_MAX.  *
 * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * */
#ifdef __FreeBSD__
	wd = malloc(MAXPATHLEN);
	if (wd == NULL) {
		syslog(LOG_ERR, "unable to allocate PATH_MAX memory: %s", strerror(errno));
		argptr->err = 0;
		goto cleanup;
	}
	if (getcwd(wd, MAXPATHLEN) == NULL) {
		argptr->err = 0;
		goto cleanup;
	}
#else
	if ((wd = getcwd(NULL, PATH_MAX)) == NULL) {
		argptr->err = 0;
		goto cleanup;
	}
#endif
	syslog(LOG_DEBUG, "cwd: %s :: realpath: %s", wd, fname);
	if (read) {
		if (strncmp(wd, fname, strlen(wd))) {
			syslog(
				LOG_ERR,
				"%s:%d requested invalid file %s",
				inet_ntoa(argptr->client.sin_addr),
				ntohs(argptr->client.sin_port),
				fname
			);
			argptr->err = 2;
			goto cleanup;
		}
	} else {
		strcat(wd, "/");
		strcat(wd, fname);
		char *intermediate = wd;
		wd = fname;
		fname = intermediate;
	}
cleanup:
	free(wd);
	free(path);
	free(parts);
	return fname;
}

ssize_t senderror(int s, nbd_nbtpd_args *argptr) {
	size_t buflen = 4 + strlen(nbd_tftp_error_to_message((nbd_tftp_ecode)argptr->err));
	char *buf = nbd_tftp_ser_error_from_code((nbd_tftp_ecode)argptr->err);
	ssize_t sb = send(s, buf, buflen, 0);
	free(buf);
	return sb;
}

void *read_req_resp(void *args) {
	char *fname = NULL, *buf = NULL, *rxb = NULL;
	FILE *fp = NULL;
	nbd_nbtpd_args *argptr = (nbd_nbtpd_args *)args;
	syslog(
		LOG_DEBUG,
		"%s:%d is requesting file %s",
		inet_ntoa(argptr->client.sin_addr),
		ntohs(argptr->client.sin_port),
		argptr->path
	);
	fname = checkpath(argptr, 1);
	if (fname == NULL) {
		goto pre_socket;
	}
	if (!is_netascii_str((char *)&(argptr->mode))) {
		syslog(
			LOG_ERR,
			"%s:%d mode is invalid %s",
			inet_ntoa(argptr->client.sin_addr),
			ntohs(argptr->client.sin_port),
			argptr->mode
		);
		argptr->err = 4;
		goto pre_socket;
	}
	nbd_opmode opmode = get_mode((char *)&(argptr->mode));
	if ((opmode != NETASCII) && (opmode != OCTET)) {
		argptr->err = 4;
		goto pre_socket;
	}
	fp = fopen(fname, "r");
	if (fp == NULL) {
		syslog(
			LOG_ERR,
			"failed to open file for reading: %s",
			strerror(errno)
		);
		argptr->err = 2;
		goto pre_socket;
	}
	buf = malloc(512);
	rxb = malloc(256);
	if ((buf == NULL) || (rxb == NULL)) {
		syslog(
			LOG_CRIT,
			"failed to allocate memory: %s",
			strerror(errno)
		);
		argptr->err = 0;
		goto pre_socket;
	}
	int s = makesock(argptr);
	if (s <= 0) {
		goto cleanup;
	}
	uint8_t rxon = 0, lon = 1;
	uint16_t bnum = 0;
	while (lon) {
		memset(buf, '\0', 512);
		uint8_t verif = 0;
		uint8_t vcount = 0;
		size_t num_read = fread(buf, 1, 512, fp);
		if (((num_read < 512) && (!feof(fp))) || (vcount > 5)) {
			argptr->err = 0;
			senderror(s, argptr);
			goto cleanup;
		}
		char *packet = nbd_tftp_ser_data_from_parts(++bnum, buf, num_read);
		while (!verif) {
			syslog(LOG_DEBUG, "sending block number %d to %s:%d",
				bnum,
				inet_ntoa(argptr->client.sin_addr),
				htons(argptr->client.sin_port)
			);
			if (send(s, packet, num_read + 4, 0) < 0) {
				syslog(
					LOG_ERR,
					"unable to send data to %s: %s",
					inet_ntoa(argptr->client.sin_addr),
					strerror(errno)
				);
				continue;
			}
			rxon = 1;
			uint8_t rxcount = 0;
			while (rxon) {
				syslog(LOG_DEBUG, "waiting for ack %d from %s:%d",
					bnum,
					inet_ntoa(argptr->client.sin_addr),
					htons(argptr->client.sin_port)
				);
				memset(rxb, '\0', 256);
				ssize_t rcv = recv(s, rxb, 256, 0);
				syslog(LOG_DEBUG, "recv'd %ld bytes from %s:%d",
					rcv,
					inet_ntoa(argptr->client.sin_addr),
					htons(argptr->client.sin_port)
				);
				if (rcv == 4) {
					nbd_tftp_packet_ack ack = nbd_tftp_de_ack(rxb, rcv);
					syslog(LOG_DEBUG, "%s:%d ack'd with block number: %d (expect %d)",
						inet_ntoa(argptr->client.sin_addr),
						htons(argptr->client.sin_port),
						ack.block_num,
						bnum
					);
					if (ack.block_num == bnum) {
						rxon = 0;
						verif = 1;
						break;
					}
					if (ack.block_num == (bnum - 1)) {
						rxon = 0;
						vcount++;
						break;
					}
				}
				if (++rxcount > 30) {
					argptr->err = 0;
					senderror(s, argptr);
					goto cleanup;
				}
			}
		}
		free(packet);
		if (num_read < 512) {
			lon = 0;
		}
	}
	close(s);
cleanup:
	fclose(fp);
	free(buf);
	free(rxb);
	free(fname);
	free(args);
	return (void *)NULL;
pre_socket:
	if (fp != NULL) {
		fclose(fp);
	}
	free(buf);
	free(rxb);
	free(fname);
	return nbd_nbtpd_resp_error(args);
}

void *write_req_resp(void *args) {
	char *fname = NULL, *rxb = NULL;
	FILE *fp = NULL;
	nbd_nbtpd_args *argptr = (nbd_nbtpd_args *)args;
	syslog(
		LOG_DEBUG,
		"%s:%d is trying to write file %s",
		inet_ntoa(argptr->client.sin_addr),
		ntohs(argptr->client.sin_port),
		argptr->path
	);
	fname = checkpath(argptr, 0);
	if (fname == NULL) {
		syslog(LOG_ERR, "fname is NULL");
		goto pre_socket;
	}
	if (!is_netascii_str((char *)&(argptr->mode))) {
		syslog(
			LOG_ERR,
			"%s:%d mode is invalid %s",
			inet_ntoa(argptr->client.sin_addr),
			ntohs(argptr->client.sin_port),
			argptr->mode
		);
		argptr->err = 4;
		goto pre_socket;
	}
	nbd_opmode opmode = get_mode((char *)&(argptr->mode));
	if ((opmode != NETASCII) && (opmode != OCTET)) {
		syslog(LOG_ERR, "%s:%d mode is not supported.", inet_ntoa(argptr->client.sin_addr), ntohs(argptr->client.sin_port));
		argptr->err = 4;
		goto pre_socket;
	}
	fp = fopen(fname, "w");
	if (fp == NULL) {
		syslog(
			LOG_ERR,
			"failed to open file for writing: %s",
			strerror(errno)
		);
		argptr->err = 2;
		goto pre_socket;
	}
	rxb = malloc(768); // 768 is 1.5x expected block size
	if (rxb == NULL) {
		syslog(
			LOG_CRIT,
			"failed to allocate memory: %s",
			strerror(errno)
		);
		argptr->err = 0;
		goto pre_socket;
	}
	int s = makesock(argptr);
	if (s <= 0) {
		goto cleanup;
	}
	uint8_t rxon = 0, lon = 1;
	uint16_t bnum = 0;
	ssize_t rcv = 516;
	while (lon) {
		uint8_t verif = 0;
		uint8_t vcount = 0;
		char *packet = nbd_tftp_ser_ack_from_block_num(bnum++);
		while (!verif) {
			syslog(LOG_DEBUG, "sending ack number %d to %s:%d",
				bnum,
				inet_ntoa(argptr->client.sin_addr),
				htons(argptr->client.sin_port)
			);
			if (send(s, packet, 4, 0) < 0) {
				syslog(
					LOG_ERR,
					"unable to send ack to %s:%d: %s",
					inet_ntoa(argptr->client.sin_addr),
					htons(argptr->client.sin_port),
					strerror(errno)
				);
				continue;
			}
			if ((rcv < 516) || (vcount > 5)) {
				lon = 0;
				break;
			}
			rxon = 1;
			uint8_t rxcount = 0;
			while (rxon) {
				syslog(LOG_DEBUG, "waiting for data from %s:%d",
					inet_ntoa(argptr->client.sin_addr),
					htons(argptr->client.sin_port)
				);
				memset(rxb, '\0', 768);
				rcv = recv(s, rxb, 768, 0);
				syslog(LOG_DEBUG, "recv'd %ld bytes from %s:%d",
					rcv,
					inet_ntoa(argptr->client.sin_addr),
					htons(argptr->client.sin_port)
				);
				if (rcv >= 4) {
					nbd_tftp_packet_data data = nbd_tftp_de_data(rxb, rcv);
					syslog(
						LOG_DEBUG,
						"%s:%d sent block number: %d (expect %d)",
						inet_ntoa(argptr->client.sin_addr),
						htons(argptr->client.sin_port),
						data.block_num,
						bnum
					);
					if (data.block_num == bnum) {
						rxon = 0;
						verif = 1;
						if (data.datalen > 0) {
							if (fwrite(data.data, 1, data.datalen, fp) < data.datalen) {
								syslog(LOG_ERR, "filewrite failed for %s: %s", fname, strerror(errno));
								argptr->err = 0;
								senderror(s, argptr);
								goto clean_socket;
							}
						}
						break;
					}
					if (data.block_num == (bnum - 1)) {
						rxon = 0;
						vcount++;
						break;
					}
				}
				if (++rxcount > 30) {
					syslog(LOG_ERR, "retry count exceeded");
					argptr->err = 0;
					senderror(s, argptr);
					goto clean_socket;
				}
			}
		}
		free(packet);
	}
clean_socket:
	close(s);
cleanup:
	fclose(fp);
	free(rxb);
	free(fname);
	free(args);
	return (void *)NULL;
pre_socket:
	if (fp != NULL) {
		fclose(fp);
	}
	free(rxb);
	free(fname);
	return nbd_nbtpd_resp_error(args);
}

void *nbd_nbtpd_resp_error(void *args) {
	nbd_nbtpd_args *argptr = (nbd_nbtpd_args *)args;
	int s = makesock(argptr);
	if (s <= 0) {
		syslog(
			LOG_ERR,
			"cannot send error to %s:%d",
			inet_ntoa(argptr->client.sin_addr),
			ntohs(argptr->client.sin_port)
		);
		goto cleanup;
	}
	senderror(s, argptr);
cleanup:
	free(args);
	return (void *)NULL;
}