#define _GNU_SOURCE		/* getopt_long */

#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <poll.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>

#define VERSION "1.1.1"

/* Size of the buffer used to clone stdin. */
#define BUFSIZE 8000

/* Clone of stdin.  Used to perform lseek(2) calls between each child
   process. */
FILE *stdin_clone = NULL;

/* File descriptor associated to the clone of stdin. */
int clone_fd;

void
print_usage(FILE *output)
{
	fprintf(output, "Usage:\n");
	fprintf(output, "rpt [-n | --count COUNT] [-f | --force] COMMAND\n");
	fprintf(output, "rpt [-h | --help]\n");
	fprintf(output, "rpt [-V | --version]\n");
}

void
print_help()
{
	puts("Repeat : repeat a shell command");
	puts("");
	print_usage(stdout);
}

void
print_version()
{
	puts(VERSION);
}

/* Exit whenever a subprocess fails. */
void
exit_on_error(int status)
{
	if (!status)
		return;

	fprintf(stderr, "rpt: child process exited with status %d, aborting\n",
		status);
	exit(status);
}

/* Ignore subprocesses' errors. */
void
continue_on_error(int status)
{
	(void) status;
}

void
exec_child(char **argv)
{
	int status;

	if (stdin_clone)
	{
		status = lseek(clone_fd, 0, L_SET);

		if (status < 0)
		{
			perror("lseek");
			exit(4);
		}
	}

	status = dup2(clone_fd, STDIN_FILENO);

	if (status < 0)
	{
		perror("dup2");
		exit(4);
	}

	status = execvp(argv[0], argv);

	if (status < 0)
	{
		perror("execvp");
		exit(4);
	}

	exit(4);
}

void
wait_child(void (*handle_error_f)(int status))
{
	int wstatus;

	wait(&wstatus);

	handle_error_f(WEXITSTATUS(wstatus));
}

void
invoke_cmd(char **argv, void (*handle_error_f)(int status))
{
	pid_t pid = fork();

	if (pid < 0)
	{
		perror("fork");
		exit(4);
	}

	if (pid == 0)
		exec_child(argv);
	else
		wait_child(handle_error_f);
}

/* Invoke the executable *ARGV COUNT times, passing it the rest of ARGV as
   arguments. Call HANDLE_ERROR_F between each process. */
void
repeat_cmd(char **argv, long count, void (*handle_error_f)(int status))
{
	for(long i = 0; i < count; ++i)
		invoke_cmd(argv, handle_error_f);
}

void
copy_stdin()
{
	ssize_t nbytes;
	char buf[BUFSIZE];

	stdin_clone = tmpfile(); /* closed at the end of the main */

	if (!stdin_clone)
	{
		perror("tmpfile");
		exit(1);
	}

	clone_fd = fileno(stdin_clone);

	while ((nbytes = read(STDIN_FILENO, buf, BUFSIZE)) > 0)
	{
		if (write(clone_fd, buf, nbytes) < 0)
			break;
	}

	if (errno)
	{
		fclose(stdin_clone);
		perror("read/write");
		exit(1);
	}
}

int
main(int argc, char* argv[])
{
	struct option const opts[] = {
		{"count", required_argument, NULL, 'n'},
		{"force", no_argument, NULL, 'f'},
		{"version", no_argument, NULL, 'V'},
		{"help", no_argument, NULL, 'h'},
		{NULL, 0, NULL, 0}
	};

	opterr = 0;		/* dismiss getopt error message */
	int opt;

	char *strcount = NULL;
	long count = 1;

	void (*handle_error_f)(int status) = exit_on_error;

	while ((opt = getopt_long(argc, argv, "+:n:fVh", opts, NULL)) >= 0)
	{
		switch(opt)
		{
		case 'n':
			strcount = optarg;
			break;
		case 'f':
			handle_error_f = continue_on_error;
			break;
		case 'h':
			print_help();
			return 0;
		case 'V':
			print_version();
			return 0;
		case ':':
			fprintf(stderr, "rpt: COUNT value missing\n");
			print_usage(stderr);
			return 1; /* invalid option */
		default:
			print_usage(stderr);
			return 1; /* invalid option */
		}
	}

	if (strcount)
	{
		char *err = NULL;
		count = strtol(strcount, &err, 10);

		if (err && (*err != '\0'))
		{
			fprintf(stderr, "rpt: failed to convert COUNT : '%s'\n",
				strcount);
			return 2;
		}
	}

	if (optind >= argc)
	{
		fprintf(stderr, "rpt: no command provided\n");
		return 3;
	}

	struct pollfd pollfd = {.fd=STDIN_FILENO, .events=POLLIN, .revents=0};

	if (poll(&pollfd, 1, 0) < 0)
	{
		perror("poll");
		return 4;
	}

	if (pollfd.revents & POLLIN)
		copy_stdin();

	repeat_cmd(argv + optind, count, handle_error_f);

	if (stdin_clone)
		fclose(stdin_clone);

	return 0;
}