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

char *supdir = "/home/marius/svcmon/svc";
char *super_path[] = { "/home/marius/svcmon/svcsupervise", NULL, NULL };

struct svc {
	char dir[MAXNAMLEN + 1];
	pid_t supervisor;

	TAILQ_ENTRY(svc) entries;
};

TAILQ_HEAD(svcrunlist, svc) services = TAILQ_HEAD_INITIALIZER(services);
TAILQ_HEAD(svcnorunlist, svc) svc_norun = TAILQ_HEAD_INITIALIZER(svc_norun);

struct svc *
find_svc_sup(pid_t p)
{
	struct svc *np;

	TAILQ_FOREACH(np, &services, entries) {
		if (np->supervisor == p)
			return np;
	}

	return NULL;
} 

struct svc *
find_run_svc_dir(const char *dir)
{
	struct svc *np;
	TAILQ_FOREACH(np, &services, entries) {
		if (strcmp(np->dir, dir) == 0)
			return np;
	}

	return NULL;
}

struct svc *
find_norun_svc_dir(const char *dir)
{
	struct svc *np;
	TAILQ_FOREACH(np, &svc_norun, entries) {
		if (strcmp(np->dir, dir) == 0)
			return np;
	}

	return NULL;
}

struct svc *
find_svc_dir(const char *dir)
{
	struct svc *tmp;

	tmp = find_run_svc_dir(dir);
	if (tmp != NULL)
		return tmp;

	tmp = find_norun_svc_dir(dir);
	return tmp;
}

void
move_norun(struct svc *service)
{
	TAILQ_REMOVE(&services, service, entries);
	TAILQ_INSERT_TAIL(&svc_norun, service, entries);
}

void
move_run(struct svc *service)
{
	TAILQ_REMOVE(&svc_norun, service, entries);
	TAILQ_INSERT_TAIL(&services, service, entries);
}

void
remove_svc_norun(struct svc *service)
{
	TAILQ_REMOVE(&svc_norun, service, entries);

	if (service->supervisor != -1)
		fprintf(stderr, "Removed service (%s) with pid (%d) for supervisor\n", service->dir, service->supervisor);

	free(service);
}

int
direxists(const char *dir)
{
	int r = 0;
	struct stat sb;
	char *path;

	if ((asprintf(&path, "%s/%s", supdir, dir)) == -1 || path == NULL)
		err(1, "asprintf()");

	if (stat(path, &sb) != -1) {
		if (S_ISDIR(sb.st_mode))
			r = 1;
	} else {
		if (errno != ENOENT)
			perror("stat()");
	}

	free(path);

	return r;
}

int
start_supervisor(struct svc *service)
{
	char *path;
	if ((asprintf(&path, "%s/%s", supdir, service->dir)) == -1 ||
	    path == NULL) {
		perror("asprintf()");
		return 0;
	}

	pid_t p = fork();
	if (p == 0) { /* Child */
		setsid();

		super_path[1] = path;

		if (execv(super_path[0], super_path) == -1)
			err(1, "execv()");
	} else if (p > 0) { /* Parent */
		service->supervisor = p;
	} else {
		perror("fork()");
	}

	free(path);

	return 1;
}

int
start_dead()
{
	struct svc *np, *tmp;
	int ret = 1;

	TAILQ_FOREACH_SAFE(np, &svc_norun, entries, tmp) {
		if (!direxists(np->dir)) {
			remove_svc_norun(np);
			continue;
		}

		if (start_supervisor(np) == 0)
			ret = 0;
		else
			move_run(np);
	}

	return ret;
}

void
add_missing_svc(const char *name)
{
	struct svc *tmp;

	tmp = find_svc_dir(name);
	if (tmp != NULL)
		return; /* We already have it */

	tmp = malloc(sizeof(struct svc));
	if (tmp == NULL)
		err(1, "malloc()");

	if (strlcpy(tmp->dir, name, (MAXNAMLEN + 1)) >= (MAXNAMLEN + 1))
		errx(1, "strlcpy()");
	
	tmp->supervisor = -1;

	TAILQ_INSERT_TAIL(&svc_norun, tmp, entries);
}

void
scan_svcdir(int dir_fd)
{
	DIR *dir;
	struct dirent *dp;

	if ((dir = fdopendir(dir_fd)) == NULL)
		err(1, "opendir()");

	while ((dp = readdir(dir)) != NULL) {
		if (dp->d_type != DT_DIR)
			continue;
		if (dp->d_name[0] == '.')
			continue;

		add_missing_svc(dp->d_name);
	}

	rewinddir(dir);

	if (fdclosedir(dir) == -1)
		err(1, "closedir()");
}

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

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

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

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

int
try_wait()
{
	int p, s, r = 0;
	for (;;) {
		struct svc *tmp;

		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;

		tmp = find_svc_sup(p);
		if (tmp == NULL)
			continue; /* XXX: Log something here? */

		tmp->supervisor = -1;
		move_norun(tmp);
	}
	
	return r;
}

int
acquire_lock()
{
	int lock_fd;
	char *lock_path;

	if (asprintf(&lock_path, "%s/lock", supdir) == -1 || lock_path == NULL)
		err(1, "asprintf()");

	if ((lock_fd = open(lock_path, O_CREAT | O_TRUNC | O_WRONLY | O_CLOEXEC, 0644)) == -1)
		err(1, "open()");

	if (flock(lock_fd, LOCK_EX | LOCK_NB) == -1) {
		if (errno == EWOULDBLOCK) {
			fprintf(stderr, "%s is locked\n", lock_path);
			exit(1);
		} else {
			err(1, "flock()");
		}
	}

	free(lock_path);

	return lock_fd;
}

void
signal_existing(int dir_fd)
{
	DIR *dir;
	struct dirent *dp;
	char *ctrl_path, *lock_path;
	int ctrl_fd, lock_fd;
	sigset_t mask;

	sigemptyset(&mask);
	sigaddset(&mask, SIGPIPE);

	dir = fdopendir(dir_fd);

	while ((dp = readdir(dir)) != NULL) {
		if (dp->d_type != DT_DIR)
			continue;
		if (dp->d_name[0] == '.')
			continue;

		if (asprintf(&ctrl_path, "%s/supervise/control", dp->d_name) == -1 || ctrl_path == NULL)
			err(1, "asprintf()");
		if (asprintf(&lock_path, "%s/supervise/lock", dp->d_name) == -1 || lock_path == NULL)
			err(1, "asprintf()");

		if ((lock_fd = open(lock_path, O_RDONLY | O_NONBLOCK | O_CLOEXEC)) == -1)
			err(1, "open()");

		if (flock(lock_fd, LOCK_EX | LOCK_NB) == -1) {
			if (errno == EWOULDBLOCK) {
				if ((ctrl_fd = open(ctrl_path, O_WRONLY | O_CLOEXEC)) == -1)
					err(1, "open()");

				if (sigprocmask(SIG_BLOCK, &mask, NULL) == -1)
					err(1, "setprocmask()");

				if (write(ctrl_fd, "x", 1) != 1) {
					if (errno != EPIPE)
						err(1, "write()");
				}

				if (sigprocmask(SIG_UNBLOCK, &mask, NULL) == -1)
					err(1, "setprocmask()");

				if (close(ctrl_fd) == -1)
					err(1, "close()");
			} else {
				err(1, "flock()");
			}
		}

		if (close(lock_fd) == -1)
			err(1, "flock()");

		free(lock_path);
		free(ctrl_path);
	}

	rewinddir(dir);

	if (fdclosedir(dir) == -1)
		err(1, "fdclosedir()");
}

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()");
}

int
main(int argc, char **argv)
{
	int kq, dir_fd, lock_fd, e, i;
	struct kevent evt[6], revt[6];
	struct procctl_reaper_kill rk;
	struct procctl_reaper_status rs;

	if (argc > 1)
		supdir = argv[1];

	if (chdir(supdir) == -1)
		err(1, "chdir()");

	lock_fd = acquire_lock();

	setup_signals();

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

	dir_fd = open(supdir, O_RDONLY | O_DIRECTORY | O_CLOEXEC);
	if (dir_fd == -1)
		err(1, "open()");

	signal_existing(dir_fd);

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

	EV_SET(&evt[0], dir_fd,  EVFILT_VNODE,  EV_ADD | EV_ENABLE, NOTE_WRITE | NOTE_EXTEND, 0, 0);
	EV_SET(&evt[1], SIGHUP,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, 0);
	EV_SET(&evt[2], SIGINT,  EVFILT_SIGNAL, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, 0);
	EV_SET(&evt[3], SIGTERM, EVFILT_SIGNAL, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, 0);
	EV_SET(&evt[4], SIGCHLD, EVFILT_SIGNAL, EV_ADD | EV_ENABLE | EV_CLEAR, 0, 0, 0);
	EV_SET(&evt[5], 1,       EVFILT_TIMER,  EV_ADD | EV_ENABLE | EV_ONESHOT, NOTE_SECONDS, 5, 0);

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

	scan_svcdir(dir_fd);
	if (start_dead() == 0) {
		if (kevent(kq, &evt[5], 1, NULL, 0, NULL) == -1)
			perror("kevent()");
	}

	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_VNODE && revt[i].ident == dir_fd) {
					scan_svcdir(dir_fd);
				} else if (revt[i].filter == EVFILT_SIGNAL) {
					if (revt[i].ident == SIGCHLD) {
						if (try_wait()) {
							if (kevent(kq, &evt[5], 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) {
					if (start_dead() == 0) {
						if (kevent(kq, &evt[5], 1, NULL, 0, NULL) == -1)
							perror("kevent()");
					}
				} 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(dir_fd) == -1)
		perror("close()");
	if (close(lock_fd) == -1)
		perror("close()");

	reset_signals();

	fprintf(stderr, "Waiting for children to exit\n");
	reap_all();

	return 0;
}