#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <err.h>
#include <string.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/event.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <sys/file.h>
#include <sys/wait.h>
#include <sys/procctl.h>

#include "status.h"

char *const run_path[] = { "./run", NULL };

void
start_proc()
{
	pid_t p = fork();
	if (p == 0) {
		execv(run_path[0], run_path);
	} else if (p > 0) {
		/* XXX: Do something apropriate here */
	} else {
		err(1, "fork()");
	}
}

int
acquire_lock()
{
	int lock_fd;

	if ((lock_fd = open("supervise/lock", O_CREAT | O_TRUNC | O_RDONLY, 0600)) == -1)
		err(1, "open()");

	if (flock(lock_fd, LOCK_EX | LOCK_NB) == -1) {
		if (errno == EWOULDBLOCK) {
			fprintf(stderr, "Could not acquire lock\n");
			exit(1);
		} else {
			err(1, "flock()");
		}
	}

	return lock_fd;
}

void
update_status()
{
	int fd, r;
	unsigned char status[18];

	if ((fd = open("supervise/status.new", O_CREAT | O_TRUNC | O_WRONLY | 0644)) == -1) {
		perror("open()");
		return;
	}

	if ((r = write(fd, status, sizeof(status))) == -1)
		perror("write()");

	if (close(fd) == -1)
		perror("close()");

	if (r < sizeof(status)) {
		fprintf(stderr, "Failed to fully write status.new\n");

		if (unlink("supervise/status.new") == -1)
			perror("unlink()");

		return;
	}

	if (rename("supervise/status.new", "supervise/status") == -1)
		perror("rename()");
}

int
handle_ctrl_command(int fd)
{
	char cmd;
	int r;

	r = read(fd, &cmd, 1);
	if (r == 0) {
		return 1;
	} else if (r == -1) {
		if (errno != EAGAIN && errno != EINTR)
			perror("read()");
		return 1;
	} else {
		switch (cmd) {
		case '\n':
			/* Ignore newline */
			break;
		case 'd':
			/* Down */
			break;
		case 'u':
			/* Up */
			break;
		case 'o':
			/* Once, up do not restart */
			break;
		case 'a':
			/* SIGALRM */
			break;
		case 'h':
			/* SIGHUP */
			break;
		case 'k':
			/* SIGKILL */
			break;
		case 't':
			/* SIGTERM */
			break;
		case 'i':
			/* SIGINT */
			break;
		case 'p':
			/* SIGSTOP */
			break;
		case 'c':
			/* SIGCONT */
			break;
		case 'x':
			return 0; /* Exit */
			break;
		default:
			fprintf(stderr, "Unknown command \"%c\"", cmd);
			break;
		}
	}

	return 1;
}

void
setup_signals()
{
	struct sigaction act;
	memset(&act, 0, sizeof(act));

	act.sa_handler = SIG_IGN;
	if (sigaction(SIGHUP, &act, NULL) == -1)
		err(1, "sigaction()");
	if (sigaction(SIGINT, &act, NULL) == -1)
		err(1, "sigaction()");
	if (sigaction(SIGTERM, &act, NULL) == -1)
		err(1, "sigaction()");
}

void
reset_signals()
{
	struct sigaction act;
	memset(&act, 0, sizeof(act));

	act.sa_handler = SIG_DFL;
	if (sigaction(SIGHUP, &act, NULL) == -1)
		err(1, "sigaction()");
	if (sigaction(SIGINT, &act, NULL) == -1)
		err(1, "sigaction()");
	if (sigaction(SIGTERM, &act, NULL) == -1)
		err(1, "sigaction()");
}

void
reap_all()
{
	int r, m = 0;

	for (;;) {
		r = waitpid(-1, NULL, WNOHANG);

		if (r == 0) {
			if (m >= 3)
				break;
			sleep(1);
			m++;
		}

		if (r == -1) {
			if (errno == EINTR)
				continue;
			if (errno != ECHILD)
				perror("waitpid()");
			break;
		}
	}
}

int
try_wait()
{
	int p, s, r = 0;
	for (;;) {
		p = waitpid(-1, &s, WNOHANG);
		if (p == 0)
			break;
		if (p == -1) {
			if (errno == EINTR) /* Impossible? */
				continue;
			if (errno != ECHILD)
				perror("waitpid()");
			break;
		}

		r = 1;
	}

	return r;
}

int
main(int argc, char **argv)
{
	int ctrl_fd, ok_fd, lock_fd, kq, e, i, wantup;
	struct stat sb;
	struct kevent evt[7], revt[7];
	struct procctl_reaper_kill rk;
	struct procctl_reaper_status rs;

	if (argc != 2)
		errx(1, "Usage: %s <dir>\n", argv[0]);

	if (chdir(argv[1]) == -1)
		err(1, "chdir()");

	if (mkdir("supervise", 0700) == -1 && errno != EEXIST)
		err(1, "mkdir()");

	lock_fd = acquire_lock();

	if (mkfifo("supervise/control", 0600) == -1 && errno != EEXIST)
		err(1, "mkfifo()");

	if (mkfifo("supervise/ok", 0600) == -1 && errno != EEXIST)
		err(1, "mkfifo()");

	if (procctl(P_PID, getpid(), PROC_REAP_ACQUIRE, NULL) == -1)
		err(1, "procctl()");

	setup_signals();

	kq = kqueue();
	if (kq == -1)
		err(1, "kqueue()");

	if ((ctrl_fd = open("supervise/control", O_RDONLY | O_NONBLOCK | O_CLOEXEC)) == -1)
		err(1, "open()");

	if ((ok_fd = open("supervise/ok", O_RDONLY | O_NONBLOCK | O_CLOEXEC)) == -1)
		err(1, "open()");

	if (stat("./down", &sb) == -1) {
		if (errno == ENOENT) {
			wantup = 1;
		} else
			err(1, "lstat()");
	} else {
		wantup = 0;
	}

	EV_SET(&evt[0], SIGHUP,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[1], SIGINT,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[2], SIGTERM, EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[3], SIGCHLD, EVFILT_SIGNAL, EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[4], ctrl_fd, EVFILT_READ,   EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[5], ok_fd,   EVFILT_READ,   EV_ADD | EV_ENABLE, 0, 0, 0);
	EV_SET(&evt[6], 1,       EVFILT_TIMER,  EV_ADD | EV_ENABLE | EV_ONESHOT, NOTE_SECONDS, 5, 0);

	if (kevent(kq, evt, 6, NULL, 0, NULL) == -1)
		err(1, "kevent()");

	if (wantup)
		start_proc();

	for (;;) {
		e = kevent(kq, NULL, 0, revt, 6, NULL);
		if (e == -1) {
			if (errno != EINTR)
				perror("kevent()");
			continue;
		} else if (e > 0) {
			for (i = 0; i < e; i++) {
				if (revt[i].filter == EVFILT_SIGNAL) {
					if (revt[i].ident == SIGCHLD) {
						if (try_wait()) {
							if (kevent(kq, &evt[6], 1, NULL, 0, NULL) == -1)
								perror("kevent()");
						}
					}
					if (revt[i].ident == SIGHUP ||
					    revt[i].ident == SIGINT ||
					    revt[i].ident == SIGTERM) {
						goto end;
					}
				} else if (revt[i].filter == EVFILT_TIMER && revt[i].ident == 1) {
					start_proc(); /* XXX: Handle some failures here */
				} else if (revt[i].filter == EVFILT_READ) {
					if (revt[i].ident == ctrl_fd) {
						if (handle_ctrl_command(ctrl_fd) == 0)
							goto end;
					}
					if (revt[i].ident == ok_fd) {
						char buf[1024];
						read(ok_fd, buf, sizeof(buf));
						/* Make sure the ok pipe doesn't fill up */
					}
				} else {
					fprintf(stderr, "Unknown event\n");
				}
			}
		}
	}

end:
	if (procctl(P_PID, getpid(), PROC_REAP_STATUS, &rs) != -1) {
		if (rs.rs_children > 0) {
			rk.rk_sig = SIGTERM;
			rk.rk_flags = REAPER_KILL_CHILDREN;
			if (procctl(P_PID, getpid(), PROC_REAP_KILL, &rk))
				perror("procctl()");
		}
	} else {
		perror("procctl()");
	}

	if (close(kq) == -1)
		perror("close()");

	if (close(ctrl_fd) == -1)
		perror("close()");

	if (close(lock_fd) == -1)
		perror("close()");

	reset_signals();

	reap_all();

	return 0;
}