/* Copyright (C) 2010,2013 G.P. Halkes
   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License version 3, as
   published by the Free Software Foundation.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <stdlib.h>
#include <errno.h>
#include <stdarg.h>
#include <unistd.h>
#include <stdio.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/wait.h>
#include <termios.h>
#include <fcntl.h>
#include <time.h>
#include <string.h>
#include <signal.h>
#include <ctype.h>
#include <limits.h>
#include <stdint.h>
#include <libgen.h>

#ifdef HAVE_SYS_SELECT_H
#include <sys/select.h>
#endif

#include "common.h"

int signal_pipe[2];
static const char controlMap[] = "@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]_^";
static struct termios saved;
static struct timeval timestamp;
static Type last_type = TYPE_NONE;


/** Alert the user of a fatal error and quit.
    @param fmt The format string for the message. See fprintf(3) for details.
    @param ... The arguments for printing.
*/
void fatal(const char *fmt, ...) {
	va_list args;

	va_start(args, fmt);
	vfprintf(stderr, fmt, args);
	va_end(args);

//FIXME: should we always do this? Standard error is not always visible...
/*
	char *fmtBuffer, *msgBuffer;

 	va_start(args, fmt);
	asprintf(&fmtBuffer, "xmessage \"%s\"", fmt, args);
	vasprintf(&msgBuffer, fmtBuffer, args);
	system(msgBuffer);
	va_end(args); */

	exit(EXIT_FAILURE);
}

void safe_write(int fd, const void *buffer, size_t size) {
	while (size) {
		ssize_t result = write(fd, buffer, size);
		if (result < 0) {
			if (errno == EINTR || errno == EAGAIN)
				continue;

			fatal("Error writing data: %s\n", strerror(errno));
		}
		buffer += result;
		size -= result;
	}
}

ssize_t safe_read(int fd, char *buffer, size_t size) {
	size_t todo = size;

	while (todo) {
		ssize_t result = read(fd, buffer, todo);
		if (result < 0) {
			if (errno == EINTR)
				continue;
			else if (errno == EAGAIN)
				return size - todo;

			fatal("Error reading data: %s\n", strerror(errno));
		}
		buffer += result;
		todo -= result;
	}
	return size;
}

void write_echo(const unsigned char *buffer, size_t size) {
	size_t write_from = 0, i;
	char ctrl[2] = { '^', 0 };

	for (i = 0; i < size; i++) {
		/* FIXME: some of the control characters should be interpreted or sent to the
		   terminal, depending on the terminal settings.
		   C-D (4) should not be printed at all
		   return (13) and newline (10) should be printed depending on the NL to CR mapping
		   tab (9) should be sent to the terminal
		   C-U, C-W seem to kill the line
			if (buffer[i] < 32 && buffer[i] != 4 && buffer[i] != 10 && buffer[i] != 9 && buffer[i] != 13 && buffer[i] != 21 && buffer[i] != 23 && buffer[i] != 28) {
		*/
		if (buffer[i] < 32 && buffer[i] != 4 && buffer[i] != 9 && buffer[i] != 10 && buffer[i] != 13) {
			if (i > write_from)
				safe_write(STDOUT_FILENO, buffer + write_from, i - write_from);
			write_from = i + 1;
			ctrl[1] = controlMap[buffer[i]];
			safe_write(STDOUT_FILENO, ctrl, 2);
		}
	}
	if (i > write_from)
		safe_write(STDOUT_FILENO, buffer + write_from, i - write_from);
}

void copy_attrs_to_terminal(int from) {
	struct termios attrs;

	if (tcgetattr(from, &attrs) < 0)
		return;
	attrs.c_lflag &= ~(ICANON | ECHO);
	tcsetattr(STDIN_FILENO, TCSADRAIN, &attrs);
}

void set_non_block(int fd, const char *descr, td_bool value) {
	int flags;
	if ((flags = fcntl(fd, F_GETFL)) < 0)
		fatal("Could not get flags for %s\n", descr);
	if (value)
		flags |= O_NONBLOCK;
	else
		flags &= ~O_NONBLOCK;
	if (fcntl(fd, F_SETFL, flags) < 0)
		fatal("Could not set flags for %s\n", descr);
}

FILE *open_output(char **name, const char *directory) {
	FILE *retval;
	mode_t oldUmask;

	/* If no name is given, create one */
	if (*name == NULL) {
		char nameBuffer[1024];
		time_t seconds;
		struct tm *local;

		seconds = time(NULL);
		local = localtime(&seconds);
		if (!local)
			fatal("Error getting localtime\n");
 		strftime(nameBuffer, 1024, "tdrecord-%Y-%m-%d@%H:%M:%S", local);
		*name = strdup(nameBuffer);
		if (*name == NULL)
			fatal("Out of memory\n");

		/* Use directory if one is supplied */
		if (directory != NULL) {
			snprintf(nameBuffer, 1024, "%s/%s", directory, *name);
			free(*name);
			*name = strdup(nameBuffer);
			if (*name == NULL)
				fatal("Out of memory\n");
		}
	}

	oldUmask = umask((mode_t) 0077);
	if ((retval = fopen(*name, "w")) == NULL)
		fatal("Could not open output %s: %s\n", *name, strerror(errno));
	umask(oldUmask);
	return retval;
}

void install_signal_handler(int sig, void (*handler)(int), const char *descr) {
	static td_bool pipe_initialised = td_false;
	struct sigaction sa;
	sigset_t sigs;

	if (!pipe_initialised) {
		if (pipe(signal_pipe) < 0)
			fatal("Could not open pipe: %s\n", strerror(errno));
		pipe_initialised = td_true;
	}

	/* Set signal handler to catch sig */
	sa.sa_handler = handler;
	sigemptyset(&sa.sa_mask);
	sa.sa_flags = 0;
	if (sigaction(sig, &sa, NULL) < 0)
		fatal("Could not set signal handler for %s: %s\n", descr, strerror(errno));

	sigemptyset(&sigs);
	sigaddset(&sigs, sig);
	if (sigprocmask(SIG_UNBLOCK, &sigs, NULL) < 0)
		fatal("Could not block %s: %s\n", descr, strerror(errno));
}

void sigwinch_handler(int param) {
	char signal_byte = SIGNAL_WINCH;
	(void) param;
	safe_write(signal_pipe[PIPE_WRITE], &signal_byte, 1);
}

void sigchld_handler(int param) {
	char signal_byte = SIGNAL_CHILD;
	(void) param;
	safe_write(signal_pipe[PIPE_WRITE], &signal_byte, 1);
}


void save_tty(void) {
	/* Get terminal attributes. */
	if (tcgetattr(STDIN_FILENO, &saved) < 0)
		fatal("Could not get terminal settings: %s\n", strerror(errno));
	atexit(reset_tty);
}

void reset_tty(void) {
	tcsetattr(STDIN_FILENO, TCSAFLUSH, &saved);
	set_non_block(STDIN_FILENO, "terminal", td_false);
}

int timevaldiff(struct timeval origin, struct timeval now) {
	return (now.tv_usec - origin.tv_usec) / 1000 + (now.tv_sec - origin.tv_sec) * 1000;
}

void init_timestamp(void) { gettimeofday(&timestamp, NULL); }

int get_elapsed(void) {
	struct timeval now;
	int retval;

	gettimeofday(&now, NULL);
	retval = timevaldiff(timestamp, now);
	timestamp = now;
	return retval;
}

void write_record(FILE *output, Type dir, int delay, int min_delay, unsigned char *buffer, ssize_t length) {
	int i;

	if (last_type != dir) {
		if (last_type == TYPE_OUTPUT)
			fputc('"', output);
		if (last_type != TYPE_NONE)
			fputc('\n', output);
		fprintf(output, "%s", dir == TYPE_OUTPUT ? "expect " : "send");
		if (dir == TYPE_OUTPUT)
			fputc('"', output);
		last_type = dir;
	}

	if (dir == TYPE_INPUT) {
		fprintf(output, " %d ", delay);
		if (min_delay > 0)
			fprintf(output, ">%d ", min_delay);
		fputc('"', output);
	}

	for (i = 0; i < length; i++) {
		if (buffer[i] == '\r' && dir == TYPE_OUTPUT) {
			/* FIXME: don't start a new string here because it may become an empty string. */
			fputs("\\r\"\n\t\"", output);
		} else if (!isprint(buffer[i])) {
			fprintf(output, "\\%03o", buffer[i]);
		} else if (buffer[i] == '"' || buffer[i] == '\\') {
			fputc('\\', output);
			fputc(buffer[i], output);
		} else {
			fputc(buffer[i], output);
		}
	}
	if (dir == TYPE_INPUT)
		fputc('"', output);
}


void copy_fd(int from, int to, Type dir, FILE *output) {
	unsigned char buffer[1024];
	ssize_t result;
	struct termios backup;
	td_bool restore = td_true;

	if (dir == TYPE_INPUT) {
		/* Temporarily disable ECHO on the terminal to prevent duplicate output.
		   Note that the "to" terminal in this case is the pty connected to the
		   client program. */
		if (tcgetattr(to, &backup) < 0 || !(backup.c_lflag & ECHO)) {
			/* If we can't read the attributes, we can't restore them at the end either.
			   And if echo is disabled anyway, we don't have to. */
			restore = td_false;
		} else {
			struct termios noecho = backup;
			noecho.c_lflag &= ~ECHO;
			tcsetattr(to, TCSANOW, &noecho);
		}
	}

repeat:
	while ((result = read(from, buffer, sizeof(buffer))) > 0) {
		/* Record copied data to file. */
		if (output != NULL)
			write_record(output, dir, dir == TYPE_INPUT ? get_elapsed() : 0, 0, buffer, result);

		/* Simulate ECHO flag on terminal. The ECHO flag has been disabled because
		   we can't get a signal when the client program changes the terminal
		   settings. This means that the first character typed will have the old ECHO
		   settings and may echo when it is not desired. However, if the client
		   program does want the ECHO flag we must now provide... */
		if (dir == TYPE_INPUT && restore)
			write_echo(buffer, result);
		safe_write(to, buffer, result);
	}
	if (result < 0) {
		if (errno == EINTR) {
			goto repeat;
		} else if (errno != EAGAIN) {
			if (dir == TYPE_INPUT)
				fatal("Error reading data from terminal or pty: %s\n", strerror(errno));
		}
	}

	if (dir == TYPE_INPUT && restore)
		tcsetattr(to, TCSANOW, &backup);
}

int interact_loop(int master, pid_t pid, FILE *output, int timeout) {
	port_pollfd fdset[3];
	struct winsize wsz;
	int status;
	int remaining_timeout = timeout;
	struct timeval start;

	fdset[0].fd = STDIN_FILENO;
	fdset[0].events = PORT_POLLIN;
	fdset[1].fd = master;
	fdset[1].events = PORT_POLLIN;
	fdset[2].fd = signal_pipe[PIPE_READ];
	fdset[2].events = PORT_POLLIN;

	copy_attrs_to_terminal(STDIN_FILENO);

	if (timeout >= 0)
		gettimeofday(&start, NULL);

	while (1) {
		int result;

		result = port_poll(fdset, 3, remaining_timeout);
		if (result == -1) {
			if (errno != EINTR)
				fatal("Error during poll: %s\n", strerror(errno));
		}

		if (result == 0)
			return -1;

		copy_attrs_to_terminal(master);

		if (fdset[1].revents & PORT_POLLIN)
			copy_fd(master, STDOUT_FILENO, TYPE_OUTPUT, output);

		copy_attrs_to_terminal(master);

		if (fdset[0].revents & PORT_POLLIN)
			copy_fd(STDIN_FILENO, master, TYPE_INPUT, output);

		/* Check if a signal handler has been called since the last check. */
		if (fdset[2].revents & PORT_POLLIN) {
			HANDLE_SIG_BYTES()
				case SIGNAL_CHILD:
					if (waitpid(pid, &status, WNOHANG | WUNTRACED) == pid) {
						if (WIFSTOPPED(status)) {
							if (output != NULL) {
								if (last_type == TYPE_OUTPUT)
									fputs("\"\n", output);
								fprintf(output, "expect_suspend\n");
								last_type = TYPE_NONE;
							}

							printf("Process has stopped. Will send continue signal in 3 seconds\n");
							sleep(3);
							kill(pid, SIGCONT);
							if (timeout >=0)
								timeout += 3000;
						} else {
							goto end;
						}
					}
					break;
				case SIGNAL_WINCH:
					if (ioctl(STDIN_FILENO, TIOCGWINSZ, &wsz) < 0)
						fatal("Error obtaining window size: %s\n", strerror(errno));
					if (output != NULL) {
						if (last_type == TYPE_OUTPUT)
							fputs("\"\n", output);
						fprintf(output, "window_size %d %d %d\n", wsz.ws_col, wsz.ws_row, get_elapsed());
						last_type = TYPE_NONE;
					}
					if (ioctl(master, TIOCSWINSZ, &wsz) < 0)
						fatal("Error setting window size: %s\n", strerror(errno));
					break;
			END_HANDLE
		}

		if (timeout >= 0) {
			struct timeval now;
			gettimeofday(&now, NULL);
			remaining_timeout = timeout - timevaldiff(start, now);
			if (remaining_timeout <= 0)
				return -1;
		}
	}

end:
	/* Copy any remaining data */
	copy_attrs_to_terminal(master);
	copy_fd(master, STDOUT_FILENO, TYPE_OUTPUT, output);
	if (output != NULL) {
		if (last_type == TYPE_OUTPUT)
			fputs("\"\n", output);
		fprintf(output, "expect_exit %d\n", WEXITSTATUS(status));
	}
	return status;
}

void write_escaped_string(FILE *out, const char *string, size_t length) {
	size_t i;

	fputc('"', out);
	for (i = 0; i < length; i++, string++) {
		if (!isprint(*string)) {
			fprintf(out, "\\%03o", (unsigned char) *string);
		} else if (*string == '\\' || *string == '"') {
			fputc('\\', out);
			fputc(*string, out);
		} else {
			fputc(*string, out);
		}
	}
	fputs("\"", out);
}

td_bool detect_resize_capable(void) {
	/* Only way to check whether the terminal is resize capable is to check the TERM
	   environment variable. Unfortunately, this is not 100% sure. For example, xterm
	   only allows this when "XTerm*allowWindowOps" resource is set to td_true. */
	const char *term = getenv("TERM");

	if (term == NULL)
		return td_false;

	return strcmp(term, "xterm") == 0 || strcmp(term, "rxvt") == 0 || strcmp(term, "Eterm") == 0;
}

pid_t start_client(ExpNode *client, int fd, FILE *log_file) {
	pid_t result;
	StringListNode *ptr;
	int argc = 0;
	char **argv;

	if (log_file != NULL) {
		/* Log the start attempt */
		fprintf(log_file, "Starting client:");
		ptr = client->un.args;
		while (ptr != NULL) {
			fprintf(log_file, " %s", ptr->string);
			ptr = ptr->next;
		}
		fprintf(log_file, "\n");
		fflush(log_file);
	}

	/* Start the client */
	if ((result = fork()) < 0)
		fatal("Error executing fork: %s\n", strerror(errno));
	else if (result != 0)
		return result;

	if (port_login_tty(fd) < 0)
		fatal("Error executing login_tty: %s\n", strerror(errno));

	ptr = client->un.args;
	while (ptr != NULL) {
		argc++;
		ptr = ptr->next;
	}

	if ((argv = malloc(sizeof(char *) * (argc + 1))) == NULL)
		fatal("Could not allocate client argument array: %s\n", strerror(errno));

	ptr = client->un.args;
	argc = 0;
	while (ptr != NULL) {
		argv[argc++] = ptr->string;
		ptr = ptr->next;
	}
	argv[argc] = NULL;

	/* Don't want program to communicate with X11 server for modifier states. */
	unsetenv("DISPLAY");

	execvp(argv[0], argv);
	fatal("execvp failed: %s\n", strerror(errno));
	return 0;
}

char *getcwd_wrapper(void) {
	size_t cwdBufferLength = 1024;
	char *cwdBuffer, *retval;

	if ((cwdBuffer = malloc(cwdBufferLength)) == NULL)
		return strdup("<memory allocation error on retrieving current working directory>");

repeat_getcwd:
	retval = getcwd(cwdBuffer, cwdBufferLength);
	if (retval != NULL)
		return retval;

	if (errno == ERANGE) {
		char *newCwdBuffer;
		if (SIZE_MAX / 2 > cwdBufferLength) {
			free(cwdBuffer);
			return strdup("<memory allocation error on retrieving current working directory>");
		}
		cwdBufferLength <<= 1;
		newCwdBuffer = realloc(cwdBuffer, cwdBufferLength);
		if (newCwdBuffer == NULL) {
			free(cwdBuffer);
			return strdup("<memory allocation error on retrieving current working directory>");
		}
		goto repeat_getcwd;
	}

	free(cwdBuffer);
	/* FIXME: report _which_ error */
	return strdup("<error on retrieving current working directory>");
}

void check_script(const char *file_name, ExpNode *script) {
	int current_check_version = 0;
	for (; script != NULL; script = script->next) {
		switch (script->type) {
			case EXP_START:
			case EXP_INTERACT:
			case EXP_ENV:
			case EXP_EXPECT:
			case EXP_EXPECT_EXIT:
			case EXP_EXPECT_SUSPEND:
				/* Ignore stuff that isn't relavant to simply checking the script. */
				break;
			case EXP_SEND: {
				StringListNode *string;

				if (current_check_version >= 2)
					break;

				for (string = script->un.args; string != NULL; string = string->next) {
					if (string->min_delay > 0) {
						fprintf(stderr, "%s:%d: WARNING: minimum delay construct is not supported in script version 1. "
						                "Please add 'require_version 2' to your script.\n", file_name, script->line);
						break;
					}
				}
				break;
			}
			case EXP_WINDOW_SIZE:
				if (script->un.size.min_delay > 0)
					fprintf(stderr, "%s:%d: WARNING: minimum delay construct is not supported in script version 1. "
									"Please add 'require_version 2' to your script.\n", file_name, script->line);

				break;
			case EXP_REQUIRE_VERSION:
				if (script->un.number > CURRENT_SCRIPT_VERSION)
					fatal("Script requires an unsupported script version.\n");
				current_check_version = script->un.number;
				break;
			default:
				fatal("Internal error\n");


		}
	}
}

#ifndef HAVE_OPENPTY
int port_openpty(int *master, int *slave, struct winsize *wsz) {
	char *slave_name;
	int save_errno;

# ifdef HAVE_POSIX_OPENPT
	*master = posix_openpt(O_RDWR | O_NOCTTY);
# else
	*master = open("/dev/ptmx", O_RDWR | O_NOCTTY);
# endif
	if (*master == -1)
		return -1;

	if (grantpt (*master) == -1 ||
		unlockpt (*master) == -1 ||
		(slave_name = ptsname(*master)) == NULL ||
		(*slave = open(slave_name, O_RDWR | O_NOCTTY)) == -1)
	{
		save_errno = errno;
		goto error_end_close_master;
	}

	if (wsz == NULL || ioctl(*master, TIOCSWINSZ, wsz) != -1)
		return 0;

	/* ioctl failed. */
	save_errno = errno;

	close(*slave);
error_end_close_master:
	close(*master);
	errno = save_errno;
	return -1;
}
#endif

#ifndef HAVE_LOGIN_TTY
int port_login_tty(int fd) {
	if (setsid() == -1)
		return -1;
	if (ioctl(fd, TIOCSCTTY, 0) == -1)
		return -1;
	if (dup2(fd, STDIN_FILENO) == -1)
		return -1;
	if (dup2(fd, STDOUT_FILENO) == -1)
		return -1;
	if (dup2(fd, STDOUT_FILENO) == -1)
		return -1;
	close(fd);
	return 0;
}
#endif

#ifndef HAVE_FORKPTY
pid_t port_forkpty(int *master, struct winsize *wsz) {
	pid_t pid;
	int slave, save_errno;

	if (port_openpty(master, &slave, wsz) == -1)
		return -1;

	if ((pid = fork()) == -1) {
		save_errno = errno;
		goto error_end_close;
	}

	if (pid != 0)
		return pid;

	if (port_login_tty(slave) == -1) {
		save_errno = errno;
		goto error_end_close;
	}
	return 0;

error_end_close:
	close(slave);
	close(*master);
	errno = save_errno;
	return -1;
}
#endif


#ifndef HAVE_POLL
int port_poll(port_pollfd *fds, port_nfds_t nfds, int timeout) {
	fd_set readfds, writefds, exceptfds;
	struct timeval _timeout;
	port_nfds_t i;
	int maxfd = 0, result;

	FD_ZERO(&readfds);
	FD_ZERO(&writefds);
	FD_ZERO(&exceptfds);
	for (i = 0; i < nfds; i++) {
		if (fds[i].events & PORT_POLLIN)
			FD_SET(fds[i].fd, &readfds);
		if (fds[i].events & PORT_POLLOUT)
			FD_SET(fds[i].fd, &writefds);
		if (fds[i].events & PORT_POLLERR)
			FD_SET(fds[i].fd, &exceptfds);
		if (fds[i].events != 0 && fds[i].fd > maxfd)
			maxfd = fds[i].fd;
	}
	_timeout.tv_sec = timeout / 1000;
	_timeout.tv_usec = (timeout % 1000) * 1000;

	result = select(maxfd + 1, &readfds, &writefds, &exceptfds, timeout < 0 ? NULL : &_timeout);
	if (result == -1)
		return -1;
	if (result == 0)
		return 0;

	result = 0;
	for (i = 0; i < nfds; i++) {
		fds[i].revents = 0;
		if (FD_ISSET(fds[i].fd, &readfds))
			fds[i].revents |= PORT_POLLIN;
		if (FD_ISSET(fds[i].fd, &writefds))
			fds[i].revents |= PORT_POLLOUT;
		if (FD_ISSET(fds[i].fd, &exceptfds))
			fds[i].revents |= PORT_POLLERR;
		if (fds[i].revents != 0)
			result++;
	}

	return result;
}
#endif

#ifndef HAVE_STRDUP
char *port_strdup(const char *s) {
	char *result = malloc(strlen(s) + 1);
	if (result == NULL)
		return NULL;
	strcpy(result, s);
	return result;
}
#endif

const char *get_script_env(const ExpNode *script, const char *name) {
	for (; script != NULL; script = script->next) {
		if (script->type == EXP_ENV && strcmp(script->un.args->string, name) == 0)
			return script->un.args->next->string;

	}
	return NULL;
}

/* Check if a format string contains exactly 1 d-type conversion. */
td_bool check_fmt(const char *str) {
	td_bool percent_seen = td_false;
	td_bool conversion_seen = td_false;
	for (; *str != 0; str++) {
		if (*str == '%') {
			if (percent_seen)
				return td_false;
			if (str[1] == '%') {
				str++;
				continue;
			}
			percent_seen = td_true;
		} else if (percent_seen) {
			if (*str == 'd') {
				percent_seen = td_false;
				if (conversion_seen)
					return td_false;
				conversion_seen = td_true;
			} else if (*str < '0' || *str > '9') {
				return td_false;
			}
		}
	}
	return !percent_seen && conversion_seen;
}

int tdrecord_main(int argc, char *argv[]);
int tdrerecord_main(int argc, char *argv[]);
int tdreplay_main(int argc, char *argv[]);
int tdview_main(int argc, char *argv[]);
int tdcompare_main(int argc, char *argv[]);

int main(int argc, char *argv[]) {
	char *name = basename(argv[0]);

	if (strcmp(name, "tdrecord") == 0)
		return tdrecord_main(argc, argv);
	if (strcmp(name, "tdreplay") == 0)
		return tdreplay_main(argc, argv);
	if (strcmp(name, "tdrerecord") == 0)
		return tdrerecord_main(argc, argv);
	if (strcmp(name, "tdview") == 0)
		return tdview_main(argc, argv);
	if (strcmp(name, "tdcompare") == 0)
		return tdcompare_main(argc, argv);
	fprintf(stderr, "Please use the tdrecord, tdreplay, tdrerecord, tdview, tdcompare symlinks instead of running this executable\n");
	return 0;
}
