/* Copyright 2001  Mark Pulford <mark@kyne.com.au>
 * This file is subject to the terms and conditions of the GNU General Public
 * License. Read the file COPYING found in this archive for details, or
 * visit http://www.gnu.org/copyleft/gpl.html
 */

#ifdef HAVE_CONFIG_H
#include <config.h>
#endif

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/wait.h>

#ifdef HAVE_GETOPT_H
#include <getopt.h>
#endif

/* Split conditions */
int split_bytes = 0;
int split_lines = 0;

int padding_size = 1;

char **command_argv;
int command_count = 0;
int fatal_code = 255;

int pipe_exec(char **argv, int write);
int split_output(int *fd, char *buf, int size);
int safe_write(int *fd, const char *buf, int size);
void write_zero(int fd, int size);
void reap_child(int *fd);
char next_arg(int argc, char **argv);
void check_args(int argc, char **argv);
int size2num(int *size, const char *s);
int str2num(int *size, const char *s);
void version();
void usage();

int main(int argc, char **argv)
{
	char *buf;
	int bsize = 4096;
	int rsize;
	int fd = -1;
	int done = 0;

	check_args(argc, argv);

	buf = malloc(bsize);
	if(!buf) {
		fprintf(stderr, "xin: Out of memory\n");
		exit(-1);
	}
	signal(SIGPIPE, SIG_IGN);
	
	/* Only exec the command if there is input that needs to be handled */
	while( (rsize = read(0, buf, bsize)) ) {
		if(rsize < 0) {
			if(errno == EINTR)
				continue;
			perror("xin: read");
			exit(-1);
		}
		done = split_output(&fd, buf, rsize);
	}
	if(fd != -1) {
		if(done % padding_size)
			write_zero(fd, padding_size - done % padding_size);
		reap_child(&fd);
	}

	return 0;
}

void write_zero(int fd, int size)
{
	char buf[4096];
	int s;

	memset(buf, 0, sizeof(buf));
	while(size > 0) {
		s = sizeof(buf)<size ? sizeof(buf) : size;
		s = write(fd, buf, s);
		if(s < 0) {
			if(errno == EINTR)
				continue;
			perror("xin: write");
			exit(-1);
		}
		size -= s;
	}
}

/* Returns the bytes outputed so far in the current pipe */
int split_output(int *fd, char *buf, int size)
{
	static int bytes_done = 0;
	static int lines_done = 0;
	int bytes_left;
	int lines_left;
	int lines;
	int ss;
	int reap;

	while(size) {
		bytes_left = split_bytes - bytes_done;
		lines_left = split_lines - lines_done;

		lines = 0;
		reap = 0;
		ss = size;
	
		/* If we are splitting by bytes set the split size
		 * to the number of remaining bytes in this section */
		if(split_bytes && bytes_left <= ss) {
			ss = bytes_left;
			reap = 1;
		}

		/* If we are splitting by lines find the last newline
		 * and set the split size (ss) base on it.
		 * If no newline is found, just output the whole buffer */
		if(split_lines) {
			int last_newline = -1;
			int i;

			for(i=0; i<ss; i++) {
				if(buf[i] == '\n') {
					last_newline = i;
					if(++lines == lines_left) {
						reap = 1;
						break;
					}
				}
			}
			if(last_newline != -1)
				ss = last_newline + 1;
		}

		safe_write(fd, buf, ss);
		bytes_done += ss;
		lines_done += lines;
		buf += ss;
		size -= ss;

		if(reap) {
			if(bytes_done % padding_size)
				write_zero(*fd, padding_size - bytes_done %
						padding_size);
			reap_child(fd);
			bytes_done = 0;
			lines_done = 0;
		}
	}

	return bytes_done;
}

/* safe_write will output the entire buffer or die trying */
int safe_write(int *fd, const char *buf, int size)
{
	int i = 0;
	int ret;

	while(i < size) {
		if(-1 == *fd)
			*fd = pipe_exec(command_argv, 1);
		ret = write(*fd, buf+i, size-i);
		if(ret < 0) {
			if(EINTR == errno)
				continue;
			perror("xin: write");
			exit(-1);
		}
		i += ret;
	}

	return i;
}

/* close the child's fd & wait for it to exit */
void reap_child(int *fd)
{
	int status;
	int ret;

	if(*fd != -1) {
		if(close(*fd) < 0) {
			perror("xin: close");
			exit(-1);
		}
		*fd = -1;
	}

	ret = waitpid(-1, &status, 0);
	if(ret < 0) {
		perror("xin: waitpid");
		exit(-1);
	}

	if(!WIFEXITED(status) || WEXITSTATUS(status) == fatal_code) {
		fprintf(stderr, "xin: fatal error from child\n");
		exit(-1);
	}
	*fd = -1;
}

/* Return the fd to the child or die trying */
int pipe_exec(char **argv, int write)
{
	int fds[2];
	int pid;

	if(pipe(fds) < 0) {
		perror("xin: pipe");
		exit(-1);
	}

	write = write?1:0;	/* Ensure write is 1 or 0 */

	command_count++;

	pid = fork();
	if(pid < 0) {
		perror("xin: fork");
		exit(-1);
	}

	if(pid) {
		/* parent */
		if(close(fds[!write]) < 0) {
			perror("xin: close");
			exit(-1);
		}
		return fds[write];
	}
	
	/* child */
	close(fds[write]);

	if(dup2(fds[!write], !write) < 0) {
		perror("xin child: dup2");
		_exit(fatal_code);
	}

	execvp(command_argv[0], argv);
	perror("xin child: exec");

	_exit(fatal_code);	/* Let the parent know we cannot continue */
}

char next_arg(int argc, char **argv)
{
	char *optstr = "b:l:f:p:hVe";
#ifdef HAVE_GETOPT_H
	struct option optslong[] = {
		{"bytes", 1, NULL, 'b'},
		{"lines", 1, NULL, 'l'},
		{"fatal-code", 1, NULL, 'f'},
		{"help", 0, NULL, 'h'},
		{"version", 0, NULL, 'V'},
		{"exec", 0, NULL, 'e'},
		{"padding", 0, NULL, 'p'},
		{NULL, 0, NULL, 0}
	};

	return getopt_long(argc, argv, optstr, optslong, NULL);
#else
	return getopt(argc, argv, optstr);
#endif
}

/* Removed {"lines-bytes", 1, NULL, 'C'} option */
void check_args(int argc, char **argv)
{
	char ch;

	if(1 == argc) {
		version();
		exit(0);
	}
	
	ch = next_arg(argc, argv);
	while(ch != -1) {
		switch(ch) {
		case 'b':
			if(!size2num(&split_bytes, optarg)) {
				fprintf(stderr, "xin: bad format for -b or --bytes: %s\n", optarg);
				exit(-1);
			}
			break;
		case 'l':
			if(!size2num(&split_lines, optarg)) {
				fprintf(stderr, "xin: bad format for -l or --lines: %s\n", optarg);
				exit(-1);
			}
			break;
		case 'p':
			if(!size2num(&padding_size, optarg)) {
				fprintf(stderr, "xin: bad format for -p or --pad: %s\n", optarg);
				exit(-1);
			}
			break;
		case 'f':
			if(!str2num(&fatal_code, optarg)) {
				fprintf(stderr, "xin: bad argument for -f or --fatal-code: %s\n", optarg);
				exit(-1);
			}
			break;
		case 'h':
			usage();
			exit(0);
		case 'V':
			version();
			exit(0);
		case 'e':
			break;
		case ':': case '?':
			/* missing parameter or unknown option */
			exit(-1);
		default:
			abort();	/* BUG: Unimplented option */
		}
		if(ch == 'e')
			break;
		ch = next_arg(argc, argv);
	}

	if(argc <= optind) {
		fprintf(stderr, "xin: no command given\n");
		exit(-1);
	}

	/* Copy argv into NULL terminated command_argv */
	command_argv = malloc((argc - optind + 1) * sizeof(argv));
	if(!command_argv) {
		fprintf(stderr, "xin: Out of memeory\n");
		exit(-1);
	}
	memcpy(command_argv, &argv[optind], (argc - optind) * sizeof(argv));
	command_argv[argc - optind] = 0;
}

/* handles b, k, m suffixes
 * returns: 0	failure
 * 	    1	success
 * Note: this doesn't handle overflow */
int size2num(int *size, const char *s)
{
	char *suff;

	*size = strtol(s, &suff, 10);
	if(suff == s)
		return 0;	/* No number found */
	if(*size < 1)
		return 0;	/* Must be positive non zero */

	if(0 == *suff)
		return 1;	/* No suffix */

	if(0 != suff[1])
		return 0;	/* suffix too long */

	switch(*suff) {
	case 'b': case 'B':
		*size *= 512;
		return 1;
	case 'k': case 'K':
		*size *= 1024;
		return 1;
	case 'm': case 'M':
		*size *= 1048576;
		return 1;
	}

	return 0;	/* Bad modifier */
}

int str2num(int *size, const char *s)
{
	char *end;
	
	*size = strtol(s, &end, 10);
	if(*s != 0 && *end == 0)
		return 1;
	else
		return 0;
}

void usage()
{
	fprintf(stderr, "Usage: xin [OPTION...] -e COMMAND [ARGS...]\n");
	fprintf(stderr, "  -b  --bytes=SIZE       Split after SIZE bytes\n");
	fprintf(stderr, "  -l  --lines=LINES      Split after LINES lines\n");
	fprintf(stderr, "  -p  --padding=SIZE     Zero extend each section to a multiple of SIZE bytes\n");
	fprintf(stderr, "  -f  --fatal-code=CODE  Fatal error CODE returned by child COMMAND\n");
	fprintf(stderr, "  -e  --exec             Remaining arguments belong to COMMAND\n");
	fprintf(stderr, "  -V  --version          Display version\n");
	fprintf(stderr, "  -h  --help             Display help\n");
}

void version()
{
	fprintf(stderr, "xin v%s\n", VERSION);
}