/*-
 * Copyright (c) 2004-2009, Jilles Tjoelker
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with
 * or without modification, are permitted provided that the
 * following conditions are met:
 *
 * 1. Redistributions of source code must retain the above
 *    copyright notice, this list of conditions and the
 *    following disclaimer.
 * 2. Redistributions in binary form must reproduce the
 *    above copyright notice, this list of conditions and
 *    the following disclaimer in the documentation and/or
 *    other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED
 * WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
 * PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY
 * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF
 * USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
 * OF SUCH DAMAGE.
 */

#include <sys/types.h>
#include <sys/event.h>
#include <sys/sysctl.h>
#include <sys/time.h>
#include <sys/tree.h>
#include <sys/wait.h>

#include <err.h>
#include <errno.h>
#include <signal.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sysexits.h>
#include <unistd.h>

struct pid {
	RB_ENTRY(pid) entry;
	pid_t pid;
};

static int
pidcmp(const struct pid *a, const struct pid *b)
{
	return (a->pid > b->pid ? 1 : a->pid < b->pid ? -1 : 0);
}

RB_HEAD(pidtree, pid);
static struct pidtree pids = RB_INITIALIZER(&pids);
RB_GENERATE_STATIC(pidtree, pid, entry, pidcmp);

static void
usage(void)
{
	fprintf(stderr, "usage: pwait [-t timeout] [-opv] pid ...\n");
	exit(EX_USAGE);
}

/*
 * pwait - wait for processes to terminate
 */
int
main(int argc, char *argv[])
{
	struct itimerval itv;
	struct kevent *e;
	struct pid k, *p;
	char *end, *s;
	double timeout;
	size_t sz;
	long pid;
	pid_t mypid;
	int i, kq, n, ndone, nleft, opt, pid_max, ret, status;
	bool oflag, pflag, tflag, verbose;

	oflag = false;
	pflag = false;
	tflag = false;
	verbose = false;
	memset(&itv, 0, sizeof(itv));

	while ((opt = getopt(argc, argv, "opt:v")) != -1) {
		switch (opt) {
		case 'o':
			oflag = true;
			break;
		case 'p':
			pflag = true;
			break;
		case 't':
			tflag = true;
			errno = 0;
			timeout = strtod(optarg, &end);
			if (end == optarg || errno == ERANGE || timeout < 0) {
				errx(EX_DATAERR, "timeout value");
			}
			switch (*end) {
			case '\0':
				break;
			case 's':
				end++;
				break;
			case 'h':
				timeout *= 60;
				/* FALLTHROUGH */
			case 'm':
				timeout *= 60;
				end++;
				break;
			default:
				errx(EX_DATAERR, "timeout unit");
			}
			if (*end != '\0') {
				errx(EX_DATAERR, "timeout unit");
			}
			if (timeout > 100000000L) {
				errx(EX_DATAERR, "timeout value");
			}
			itv.it_value.tv_sec = (time_t)timeout;
			timeout -= (time_t)timeout;
			itv.it_value.tv_usec =
			    (suseconds_t)(timeout * 1000000UL);
			break;
		case 'v':
			verbose = true;
			break;
		default:
			usage();
			/* NOTREACHED */
		}
	}

	argc -= optind;
	argv += optind;

	if (argc == 0) {
		usage();
	}

	if ((kq = kqueue()) < 0)
		err(EX_OSERR, "kqueue");

	sz = sizeof(pid_max);
	if (sysctlbyname("kern.pid_max", &pid_max, &sz, NULL, 0) != 0) {
		pid_max = 99999;
	}
	if ((e = malloc((argc + tflag) * sizeof(*e))) == NULL) {
		err(EX_OSERR, "malloc");
	}
	ndone = nleft = 0;
	mypid = getpid();
	for (n = 0; n < argc; n++) {
		s = argv[n];
		/* Undocumented Solaris compat */
		if (strncmp(s, "/proc/", 6) == 0) {
			s += 6;
		}
		errno = 0;
		pid = strtol(s, &end, 10);
		if (pid < 0 || pid > pid_max || *end != '\0' || errno != 0) {
			warnx("%s: bad process id", s);
			continue;
		}
		if (pid == mypid) {
			warnx("%s: skipping my own pid", s);
			continue;
		}
		if ((p = malloc(sizeof(*p))) == NULL) {
			err(EX_OSERR, NULL);
		}
		p->pid = pid;
		if (RB_INSERT(pidtree, &pids, p) != NULL) {
			/* Duplicate. */
			free(p);
			continue;
		}
		EV_SET(e + nleft, pid, EVFILT_PROC, EV_ADD, NOTE_EXIT, 0, NULL);
		if (kevent(kq, e + nleft, 1, NULL, 0, NULL) == -1) {
			if (errno != ESRCH)
				err(EX_OSERR, "kevent()");
			warn("%ld", pid);
			RB_REMOVE(pidtree, &pids, p);
			free(p);
			ndone++;
		} else {
			nleft++;
		}
	}

	if ((ndone == 0 || !oflag) && nleft > 0 && tflag) {
		/*
		 * Explicitly detect SIGALRM so that an exit status of 124
		 * can be returned rather than 142.
		 */
		EV_SET(e + nleft, SIGALRM, EVFILT_SIGNAL, EV_ADD, 0, 0, NULL);
		if (kevent(kq, e + nleft, 1, NULL, 0, NULL) == -1) {
			err(EX_OSERR, "kevent");
		}
		/* Ignore SIGALRM to not interrupt kevent(2). */
		signal(SIGALRM, SIG_IGN);
		if (setitimer(ITIMER_REAL, &itv, NULL) == -1) {
			err(EX_OSERR, "setitimer");
		}
	}
	ret = EX_OK;
	while ((ndone == 0 || !oflag) && ret == EX_OK && nleft > 0) {
		n = kevent(kq, NULL, 0, e, nleft + tflag, NULL);
		if (n == -1) {
			err(EX_OSERR, "kevent");
		}
		for (i = 0; i < n; i++) {
			if (e[i].filter == EVFILT_SIGNAL) {
				if (verbose) {
					printf("timeout\n");
				}
				ret = 124;
			}
			pid = e[i].ident;
			if (verbose) {
				status = e[i].data;
				if (WIFEXITED(status)) {
					printf("%ld: exited with status %d.\n",
					    pid, WEXITSTATUS(status));
				} else if (WIFSIGNALED(status)) {
					printf("%ld: killed by signal %d.\n",
					    pid, WTERMSIG(status));
				} else {
					printf("%ld: terminated.\n", pid);
				}
			}
			k.pid = pid;
			if ((p = RB_FIND(pidtree, &pids, &k)) != NULL) {
				RB_REMOVE(pidtree, &pids, p);
				free(p);
				ndone++;
			}
			--nleft;
		}
	}
	if (pflag) {
		RB_FOREACH(p, pidtree, &pids) {
			printf("%d\n", p->pid);
		}
	}
	exit(ret);
}