/*
 * bushi.csk at gmail
 *
 * gcc -o logger logger-full.c -Wall -pthread -lutil
 *
 * ./logger XYZ > xyz.out.txt
 * ./logger -e XYZ > xyz.out.txt
 *
 */
#define _XOPEN_SOURCE   (500)     /* SUSv5 */
#define _POSIX_C_SOURCE (200809L) /* IEEE 1003.1-2008 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>

#include <pty.h>
#include <termios.h>
#include <fcntl.h>

#include <sys/select.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <signal.h>

static volatile pid_t cid;

static void termio_noecho(struct termios *tio)
{
	tio->c_lflag &= ~(ECHO | ECHONL);
}

static void termio_semiblock(struct termios *tio)
{
	tio->c_lflag &= ~(ICANON);
	tio->c_cc[VMIN] = 1;
	tio->c_cc[VTIME] = 1;
}

static int __tty_set_noecho_semiblock(int fd, int noecho, int semiblock)
{
	struct termios tios;
	int ret;

	if (!isatty(fd))
		return 0;

	ret = tcgetattr(fd, &tios);
	if (ret) {
		perror(__func__);
		return ret;
	}

	if (noecho)
		termio_noecho(&tios);
	if (semiblock)
		termio_semiblock(&tios);

	ret = tcsetattr(fd, TCSAFLUSH, &tios);
	if (ret)
		perror(__func__);
	return ret;
}

static void *recv_from_child(void *arg)
{
	int ptm_fd = *(int*)arg;
	int ret;
	char c;

	while ((ret = read(ptm_fd, &c, 1)) > 0)
		write(STDOUT_FILENO, &c, 1);

	return NULL;
}

static void recover_stdin(void *arg)
{
	struct termios *tiosp = (struct termios *)arg;
	if (tiosp)
		tcsetattr(STDIN_FILENO, TCSAFLUSH, tiosp);
}

static void *send_to_child(void *arg)
{
	struct termios tios, *tiosp = NULL;
	int ptm_fd = *(int*)arg;
	int ret, oldstate = PTHREAD_CANCEL_ENABLE;
	unsigned char c;

	ret = pthread_setcancelstate(PTHREAD_CANCEL_DISABLE, &oldstate);

	if (tcgetattr(STDIN_FILENO, &tios) == 0)
		tiosp = &tios;

	/* turn off echo, change to noncanonical mode.*/
	__tty_set_noecho_semiblock(STDIN_FILENO, 1, 1);
	setvbuf(stdin, NULL, _IONBF, 0);

	pthread_cleanup_push(recover_stdin, tiosp);

	if (ret == 0)
		pthread_setcancelstate(oldstate, NULL);

	while ((ret = read(STDIN_FILENO, &c, 1)) > 0)
		if (write(ptm_fd, &c, 1) < 0)
			break;

	pthread_cleanup_pop(1);

	if (ret < 0 && cid > 0)
		kill(cid, SIGPIPE);

	return NULL;
}

static void sig_handler(int sig)
{
	/* asynchronous-signal-safe, IEEE Std 1003.1:2013 */
	if (cid > 0)
		kill(cid, sig);
}

static int parent(int ptm_fd, int bidirect)
{
	int ret;

	/* Just redirect signals to the child to avoid zombie. */
	signal(SIGINT, sig_handler);
	signal(SIGTERM, sig_handler);

	pthread_t tid = (pthread_t)0, tid_w = (pthread_t)0;

	/* create the recv-from-child thread */
	ret = pthread_create(&tid, NULL, recv_from_child, &ptm_fd);
	if (ret) {
		fprintf(stderr, "pthread_create: %s\n", strerror(ret));
		/* avoid zombie */
		kill(cid, SIGKILL);
		waitpid(cid, NULL, 0);
		return ret;
	}

	/* create the send-to-child thread */
	if (bidirect > 0) {
		ret = pthread_create(&tid_w, NULL, send_to_child, &ptm_fd);
		if (ret) {
			pthread_kill(tid, SIGKILL);
			fprintf(stderr, "pthread_create: %s\n", strerror(ret));
			/* avoid zombie */
			kill(cid, SIGKILL);
			waitpid(cid, NULL, 0);
			return ret;
		}
	}
	
	/* wait and get child's final exit status */
	int status;
	ret = waitpid(cid, &status, 0);
	cid = 0;

	/* wait the recv-from-child thread which would exit by pipe error */
	if (tid != (pthread_t)0)
		pthread_join(tid, NULL);

	/* kill the send-to-child thread */
	if (tid_w != (pthread_t)0) {
		pthread_cancel(tid_w);
		pthread_join(tid_w, NULL);
	}

	/* print out child's final exit status */
	printf("\n");
	if (ret < 0)
		perror("waitpid");
	else if (WIFEXITED(status))
		printf("Terminated and returned %d.\n", WEXITSTATUS(status));
	else if (WIFSIGNALED(status))
		printf("Terminated by \"%s\"\n", strsignal(WTERMSIG(status)));
	else if (WIFSTOPPED(status))
		printf("Stopped by \"%s\"\n", strsignal(WSTOPSIG(status)));
	else if (WIFCONTINUED(status))
		printf("Continued by \"%s\"\n", strsignal(SIGCONT));
	else
		printf("Exited with unknown reason\n");
	fflush(stdout);
	fflush(stderr);

	return 0;
}

static void child(int ptm_fd, int argc, char **argv)
{
	char force_lang[] = "LANG=C";
	char my_marker[] = "LOGGER=1";

	argv++;
	argc--;

	if (argc > 1 && strcmp(argv[0], "-e") == 0)
		argv++;
	else
		putenv(force_lang);

	putenv(my_marker);

	int ret = execvp(argv[0], argv);

	/* never reach here. */
	perror(argv[0]);
	exit(ret);
}

int main(int argc, char **argv)
{
	if (argc < 2)
		return 1;

	int ptm_fd;
	int ret;

	cid = forkpty(&ptm_fd, NULL, NULL, NULL);

	switch (cid) {
	case -1: perror("forkpty"); break;
	case 0: child(ptm_fd, argc, argv); break;
	default: ret = parent(ptm_fd, 1); break;
	}

	close(ptm_fd);
	return ret;
}