/*
 * Copyright (C) 2024  Tristan Riehs
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "config.h"

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

/* Size of the buffer used to clone stdin. */
#define BUFSIZE (1 << 12)

/* Exit status */
#define EXIT_SUCCESS 0
#define EXIT_INVALID_ARG 1
#define EXIT_INVALID_COUNT 2
#define EXIT_NO_COMMAND 3
#define EXIT_OTHER 4

/* 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;

/* Number of times the command has been executed. */
long exec_count = 0;

/* Non-zero if exec_count has to be printed. */
int do_print_exec_count = 0;

void
print_usage(FILE *output)
{
	fprintf(output, "Usage: rpt [OPTIONS] COMMAND\n");
}

void
print_help()
{
	printf("Repeat: repeat a shell command\n");
	print_usage(stdout);
	printf("OPTIONS:\n");
	printf("  -n COUNT\tRepeat COMMAND COUNT times.\n");
	printf("  -f\t\tDo not exit if one process fails.\n");
	printf("  -h\t\tPrint this message and exit.\n");
	printf("  -v\t\tPrint the version number and exit.\n");
	printf("  -u\t\tRun COMMAND until it fails.\n");
	printf("  -p\t\tPrint how many times COMMAND has been executed.\n");
	printf("Reapeat is licensed under the GPL3 license.\n");
}

void
print_version()
{
	puts(VERSION);
}

void
print_exec_count(void)
{
	if (do_print_exec_count)
		printf("Command executed %ld time(s).\n", exec_count);
}

/* 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);

	if (stdin_clone)
		fclose(stdin_clone);

	print_exec_count();

	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, SEEK_SET);

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

	status = dup2(clone_fd, STDIN_FILENO);

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

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

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

	exit(EXIT_OTHER);
}

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

	wait(&wstatus);
	exec_count++;

	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(EXIT_OTHER);
	}

	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(EXIT_OTHER);
	}

	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(EXIT_OTHER);
	}
}

void
check_u_f_conflict(void (*handle_error_f)(int status), long count)
{
	if ((handle_error_f == continue_on_error) && (count == LONG_MAX))
	{
		fprintf(stderr, "rpt: -u and -f options both given, conflict, \
see man rpt(1)\n");
		exit(EXIT_INVALID_ARG);
	}
}

int
main(int argc, char* argv[])
{
	int opt;
	char *strcount = NULL;
	long count = 1;

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

	while ((opt = getopt(argc, argv, ":n:fuphV")) != -1)
	{
		switch(opt)
		{
		case 'n':
			strcount = optarg;
			break;
		case 'f':
			handle_error_f = continue_on_error;
			break;
		case 'u':
			count = LONG_MAX;
			break;
		case 'p':
			do_print_exec_count = 1;
			break;
		case 'h':
			print_help();
			exit(EXIT_SUCCESS);
		case 'V':
			print_version();
			exit(EXIT_SUCCESS);
		case ':':
			fprintf(stderr, "rpt: COUNT value missing\n");
			print_usage(stderr);
			exit(EXIT_INVALID_ARG);
		default:
			print_usage(stderr);
			exit(EXIT_INVALID_ARG);
		}
	}

	check_u_f_conflict(handle_error_f, count);

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

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

	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");
		exit(EXIT_OTHER);
	}

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

	repeat_cmd(argv + optind, count, handle_error_f);

	if (stdin_clone)
		fclose(stdin_clone);

	print_exec_count();

	return EXIT_SUCCESS;
}