#include <sys/types.h>
#include <sys/event.h>
#include <assert.h>
#include <err.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

static pid_t master;
static int kq;

static void
watch(uintptr_t ident, short filter, u_short flags, u_int fflags,
    intptr_t data, void *udata)
{
	struct kevent ev;

	EV_SET(&ev, ident, filter, flags, fflags, data, udata);
	if (kevent(kq, &ev, 1, NULL, 0, NULL) < 0)
		err(1, "kevent");
}

static void
dump_fflags(u_int fflags)
{
	int pipe;

	assert(fflags != 0);
	pipe = 0;
#define DUMP_FLAG(FLAG) do {						\
		if (fflags & FLAG) {					\
			printf("%s" #FLAG, pipe ? " | " : "");		\
			pipe = 1;					\
		}							\
	} while (0)

	DUMP_FLAG(NOTE_EXIT);
	DUMP_FLAG(NOTE_FORK);
	DUMP_FLAG(NOTE_EXEC);
	DUMP_FLAG(NOTE_TRACK);
	DUMP_FLAG(NOTE_TRACKERR);
	DUMP_FLAG(NOTE_CHILD);

	fflags &= ~(NOTE_EXIT | NOTE_FORK | NOTE_EXEC | NOTE_TRACK |
	    NOTE_TRACKERR | NOTE_CHILD);
	if (fflags != 0)
		printf("%s%u", pipe ? " | " : "", fflags);
}

static void
dump_event(struct kevent *ev)
{

	assert(ev->filter == EVFILT_PROC);
	printf("pid: %5d%s flags: ", (int)ev->ident,
	    ev->flags & EV_EOF ? " EV_EOF" : "");
	dump_fflags(ev->fflags);
	if (ev->data != 0)
		printf(" data: %jd", (uintmax_t)ev->data);
	printf("\n");
}

static void
child(int fd)
{
	pid_t pid;
	char c;

	if (fd > 0)
		(void)read(fd, &c, sizeof(c));
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	usleep(5000);
	exit(1);
}

static void
waitfor(int count, const char *msg)
{
	struct timespec ts;
	struct kevent ev;
	int rv;

	printf("%s:\n", msg);

	/* Wait up to 250 ms before timing out. */
	ts.tv_sec = 0;
	ts.tv_nsec = 250 * 1000 * 1000;
	for (;;) {
		rv = kevent(kq, NULL, 0, &ev, 1, &ts);
		if (rv < 0)
			err(1, "kevent");
		if (rv == 0)
			break;
		dump_event(&ev);
		--count;
	}

	if (count > 0)
		warnx("%d events missing for %s", count, msg);
	else if (count < 0)
		warnx("%d extra events for %s", -count, msg);
}

int
main(int ac, char **av)
{
	pid_t pid;
	int fds[2];
	char c;

	kq = kqueue();
	if (kq < 0)
		err(1, "kqueue");
	if (pipe(fds) < 0)
		err(1, "pipe");
	master = getpid();
	printf("master: %d\n", (int)master);

	/* Simple fork case. */
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0)
		child(fds[1]);
	watch(pid, EVFILT_PROC, EV_ADD, NOTE_FORK | NOTE_EXIT, 0, 0);
	write(fds[0], &c, sizeof(c));

	/* Should get two events, FORK and then EXIT. */
	waitfor(2, "simple fork");

	/* Test NOTE_TRACK | NOTE_EXIT. */
	watch(master, EVFILT_PROC, EV_ADD, NOTE_EXIT | NOTE_TRACK, 0, 0);
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0)
		child(-1);

	/* Should get four events: two NOTE_CHILD and two NOTE_EXIT. */
	waitfor(4, "NOTE_TRACK | NOTE_EXIT");

	/* Test NOTE_TRACK with everything. */
	watch(master, EVFILT_PROC, EV_ADD, NOTE_FORK | NOTE_EXEC | NOTE_EXIT |
	    NOTE_TRACK, 0, 0);
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0)
		child(-1);

	/*
	 * Should get six events: two each of NOTE_FORK, NOTE_CHILD,
	 * and NOTE_EXIT.
	 */
	waitfor(6, "TRACK | FORK | EXIT");

	/* Test NOTE_EXEC. */
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0) {
		usleep(5000);
		execlp("sleep", "sleep", "0.005", (char *)0);
		err(1, "execlp(\"sleep 0.005\")");
	}

	/* Should get four events: FORK, CHILD, EXEC, EXIT. */
	waitfor(4, "simple exec");

	/* Fast exec. */
	pid = fork();
	if (pid == -1)
		err(1, "fork");
	if (pid == 0) {
		execlp("true", "true", (char *)0);
		err(1, "execlp(\"true\")");
	}

	/* Up to four events, but probably fewer. */
	waitfor(4, "fast exec");
	return (0);
}