#define _XOPEN_SOURCE 600
#define _DEFAULT_SOURCE

#include <sys/select.h>
#include <sys/wait.h>
#include <sys/epoll.h>
#include <fcntl.h>
#include <termios.h>
#include <unistd.h>
#include <pty.h>
#include <errno.h>

#include <err.h>
#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <stdlib.h>

static char buf[BUFSIZ] = { 0 };

static int
writeall(int fd, char *buffer, size_t size)
{
	ssize_t nbytes;

	while (size) {
		nbytes = write(fd, buffer, size);
		if (nbytes <= 0) return 1;
		buffer += nbytes;
		size -= nbytes;
	}

	return 0;
}

static int
fwd(int in, int out)
{
	static char *buffer = NULL;
	ssize_t got, i, last;

	if (!buffer) {
		buffer = malloc(BUFSIZ);
		if (!buffer) err(1, "malloc");
	}

	do {
		got = read(in, buffer, sizeof(buffer));
		if (got < 0) return 1;

		if (writeall(out, buffer, got))
			return 1;
	} while (got == sizeof(buffer));

	return 0;
}

static void
mitm(int pts)
{
	fd_set fds;

	while (1) {
		FD_ZERO(&fds);
		FD_SET(0, &fds);
		FD_SET(pts, &fds);

		if (select(pts + 1, &fds, 0, 0, 0) == -1) {
			if (errno == EINTR) continue;
			return;
		}

		if (FD_ISSET(0, &fds)) { /* stdin -> ptty */
			if (fwd(0, pts)) return;
		}

		if (FD_ISSET(pts, &fds)) { /* ptty -> stdout + pipe */
			if (fwd(pts, 1)) return;
		}
	}
}

static void
usage(int rc)
{

	fprintf(stderr, "Usage: apty [-r ROWS] [-c COLS] [-h] [--] CMD\n");
	exit(rc);
}

int
main(int argc, char **argv)
{
	char **arg, **cmd;
	struct termios termios = {0};
	struct winsize winsize = {0};
	int exitcode;
	int master, tty;
	pid_t child;

	if (argc <= 1) usage(1);

	winsize.ws_col = 120;
	winsize.ws_row = 80;
	cfmakeraw(&termios);

	if (isatty(STDIN_FILENO)) {
		if (ioctl(STDIN_FILENO, TIOCGWINSZ, &winsize))
			err(1, "tcgetwinsize");

		if (tcgetattr(STDIN_FILENO, &termios))
			err(1, "tcgettattr");
		termios.c_lflag &= ~ICANON;
		if (tcsetattr(STDIN_FILENO, TCSANOW, &termios))
			err(1, "tcsetattr");
	}

	for (arg = argv + 1; *arg; arg++) {
		if (!strcmp(*arg, "-t") || !strcmp(*arg, "--tty")) {
			if (!isatty(STDIN_FILENO)) {
				tty = open("/dev/tty", O_RDONLY);
				if (tty < 0) err(1, "open /dev/tty");

				if (tcgetattr(tty, &termios))
					err(1, "tcgetattr tty");

				if (ioctl(tty, TIOCGWINSZ, &winsize))
					err(1, "tcgetwinsize tty");

				close(tty);
			}
		} else if (!strcmp(*arg, "-c") || !strcmp(*arg, "--cols")) {
			winsize.ws_col = atoi(*++arg);
		} else if (!strcmp(*arg, "-r") || !strcmp(*arg, "--rows")) {
			winsize.ws_row = atoi(*++arg);
		} else if (!strcmp(*arg, "-h") || !strcmp(*arg, "--help")) {
			usage(0);
		} else if (!strcmp(*arg, "--")) {
			arg++;
			break;
		} else {
			break;
		}
	}

	if (!*arg) errx(1, "missing command");
	cmd = arg;

	child = forkpty(&master, NULL, &termios, &winsize);
	if (child < 0) err(1, "fork");

	if (child == 0) {
		execvp(*cmd, cmd);
		err(1, "execvp %s", *cmd);
	} else {
		mitm(master);

		waitpid(child, &exitcode, 0);
		close(master);
	}

	return exitcode;
}