/*
 * Copyright (c) 2024 DTB <trinity@trinity.moe>
 * Copyright (c) 2024 Emma Tebibyte <emma@tebibyte.media>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 *
 * This program is free software: you can redistribute it and/or modify it under
 * the terms of the GNU Affero 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 Affero General Public License for more
 * details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program. If not, see https://www.gnu.org/licenses/.
 */

#include <assert.h> /* assert(3) */
#include <errno.h> /* errno */
#include <fcntl.h> /* open(2) */
#include <stdbool.h> /* bool */
#include <stdio.h> /* fprintf(3), stderr */
#include <stdlib.h> /* malloc(3), strtol(3), size_t */
#include <string.h> /* memcpy(3), memmove(3), memset(3) */
#include <sysexits.h> /* EX_OK, EX_OSERR, EX_USAGE */
#include <unistd.h> /* close(2), getopt(3), lseek(2), read(2), write(2),
                     * pledge(2), unveil(2), optarg, optind, STDIN_FILENO,
					 * STDOUT_FILENO */
#include <sys/stat.h> /* S_IRGRP, S_IROTH, S_IRUSR, S_IWGRP, S_IWOTH, S_IWUSR */

char *program_name = "dj";

/* dj uses two structures that respectively correspond to the reading and
 * writing ends of its jockeyed "pipe". User-configurable members are noted
 * with their relevant options. */
struct Io {
	char *buf;      /* buffer */
	char *fn;       /* file name (-io) */
	size_t bs;    	/* buffer size (-bB) */
	size_t bufuse;	/* buffer usage */
	size_t bytes; 	/* bytes processed */
	size_t prec;  	/* partial records processed */
	size_t rec;   	/* records processed */
	long seek;  	/* remaining bytes to seek/skip (-sS) */
	int error;  	/* errno */
	int fd;         /* file descriptor */
	int fl;         /* file opening flags */
};

static char *stdin_name = "<stdin>";
static char *stdout_name = "<stdout>";

static int creat_mode = S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH
	| S_IWOTH; /* Consistent with touch(1p). */
static int read_flags = O_RDONLY; /* Consistent with Busybox dd(1). */
static int write_flags = O_WRONLY | O_CREAT;

#define MIN(a, b) (((a) < (b)) ? (a) : (b))

/* Macro to check if fd is stdin or stdout */
#define fdisstd(fd) ((fd) == STDIN_FILENO || (fd) == STDOUT_FILENO)

/* Completes one Io block read */
static struct Io *
Io_read(struct Io *io) {
	int t;

	assert(io->bs > 0);
	assert(io->bufuse < io->bs);

	if ((t = read(io->fd, &(io->buf)[io->bufuse], io->bs - io->bufuse)) < 0) {
		io->error = errno;
		t = 0;
	}

	io->bufuse += t;
	io->bytes  += t;
	io->prec   += (0 < io->bufuse && io->bufuse < io->bs);
	io->rec    += (io->bufuse == io->bs);

	assert(io->bufuse <= io->bs);

	return io;
}

/* Completes one Io block write */
static struct Io *
Io_write(struct Io *io) {
	int t;

	assert(io->bufuse > 0);
	assert(io->bufuse <= io->bs);

	if ((t = write(io->fd, io->buf, io->bufuse)) < 0) {
		io->error = errno;
		t = 0;
	} else if (t > 0) {
		(void)memmove(io->buf, &(io->buf)[t], (io->bufuse -= t));
	}

	io->bytes += t;
	io->prec  += (t > 0 && io->bufuse > 0);
	io->rec   += (t > 0 && io->bufuse == 0);

	return io;
}

static int
oserr(char *e, int n) { /* program_name: [failing component:] error */
	(void)fprintf(stderr, "%s: ", program_name);
	if (e != NULL) { (void)fprintf(stderr, "%s: ", e); }
	(void)fprintf(stderr, "%s\n", strerror(n));

	return EX_OSERR;
}

/* Prints statistics regarding the use of dj, particularly partially and
 * completely read and written records. */
static int
fprintio(FILE *stream, char *fmt, struct Io io[2]) {
	return fprintf(
		stream,
		fmt,
		io[0].rec,
		io[0].prec,
		io[1].rec,
		io[1].prec,
		io[0].bytes,
		io[1].bytes
	);
}

/* To be assigned to main:fmt and used with printio(). */
static char *fmt_asv =
	"%d" /* io[0].rec   */ "\037" /* ASCII US */
	"%d" /* io[0].prec  */ "\036" /* ASCII RS */
	"%d" /* io[1].rec   */ "\037" /* ASCII US */
	"%d" /* io[1].prec  */ "\035" /* ASCII GS */
	"%d" /* io[0].bytes */ "\036" /* ASCII RS */
	"%d" /* io[1].bytes */ "\034" /* ASCII FS */
	"\n"
;
static char *fmt_human = "%d+%d > %d+%d; %d > %d\n";

/* Parses the string s to an integer, returning either the integer or in the
 * case of an error a negative integer. This is used for argument parsing
 * (e.g. -B [int]) in dj and no negative integer would be valid anyway. */
static long
parse(char *s) {
	long r;

	errno = 0;
	r = strtol(s, &s, 0);
	return (*s == '\0' /* no chars left unparsed */ && errno == 0) ? r : -1;
}

static int
usage(char *argv0) {
	(void)fprintf(
		stderr,
		"Usage: %s [-Hn] [-a byte] [-c count]\n"
		"\t[-i file] [-b block_size] [-s offset]\n"
		"\t[-o file] [-B block_size] [-S offset]\n",
		argv0
	);

	return EX_USAGE;
}

int main(int argc, char *argv[]) {
   	int align;    /* low 8b used, negative if no alignment is being done */
	int count;    /* -1 if dj(1) runs until no more reads are possible */
	char *fmt;    /* set to fmt_asv (default) or fmt_human (-H) */
	size_t i;     /* side of io (in or out) being modified */
	bool retry;   /* false if exits on partial reads or writes */
	struct Io io[2 /* { in, out } */];

#ifdef __OpenBSD__
	if (pledge("cpath rpath stdio unveil wpath", NULL) == -1) {
		return oserr(NULL, errno);
	}
#endif

	/* Set defaults. */
	align = -1;
	count = -1;
	fmt = fmt_asv;
	retry = 0;
	for (i = 0; i < (sizeof io) / (sizeof *io); ++i) {
		io[i].bs = 1024 /* 1 KiB */; /* GNU dd(1) default; POSIX says 512B */
		io[i].bufuse = 0;
		io[i].bytes = 0;
		io[i].fd = i == 0 ? STDIN_FILENO : STDOUT_FILENO;
		io[i].fn = i == 0 ? stdin_name   : stdout_name;
		io[i].fl = i == 0 ? read_flags   : write_flags;
		io[i].error = 0;
		io[i].prec = 0;
		io[i].rec = 0;
		io[i].seek = 0;
	}

	if (argc > 0) {
		int c;

		program_name = argv[0];
		while ((c = getopt(argc, argv, "a:b:B:c:i:hHns:S:o:")) != -1) {
			switch (c) {
				case 'i': case 'o': /* input, output */
					i = (c == 'o');

					/* optarg == "-" (stdin/stdout) */
					if (optarg[0] == '-' && optarg[1] == '\0') { 
						io[i].fd = i == 0 ? STDIN_FILENO : STDOUT_FILENO;
						io[i].fn = i == 0 ? stdin_name   : stdout_name;
						break;
					} else {
						int fd;

#ifdef __OpenBSD__
						if (unveil(optarg, i == 0 ? "r" : "wc") == -1) {
							return oserr(NULL, errno);
						}
#endif

						if (
							(fd = open(optarg, io[i].fl, creat_mode)) != -1
							&& (fdisstd(io[i].fd) || close(io[i].fd) == 0)
						) {
							io[i].fd = fd;
							io[i].fn = optarg;
							break;
						}
					}

					return oserr(optarg, errno);
					/* UNREACHABLE */
				case 'n': retry = 1; break; /* retry failed reads once */
				case 'H': fmt = fmt_human; break; /* human-readable output */
				case 'a': /* input buffer padding */
					if (optarg[0] == '\0' || optarg[1] == '\0') {
						align = optarg[0];
						break;
					}
					/* FALLTHROUGH */
				case 'c': /* number of reads */
				case 'b': case 'B': /* input/output block size */
				case 's': case 'S': /* (s)kip/(S)eek in input/output */
					if (c == 'c' && (count = parse(optarg)) >= 0) { break; }

					i = (c >= 'A' && c <= 'Z');
					c |= 0x20; /* 0b 0010 0000 (ASCII) make lowercase */

					if ( /* if -b or -s is parsed out correctly */
						(c == 'b' && (io[i].bs = parse(optarg)) > 0)
						|| (c == 's' && (io[i].seek = parse(optarg)) >= 0)
					) { break; } /* don't error */

					/* FALLTHROUGH */
				default:
					return usage(program_name);
			}
		}
	}

#ifdef __OpenBSD__
	if (unveil(NULL, NULL) == -1) { return oserr(NULL, errno); }
#endif

	assert(io->fd != STDIN_FILENO  || io->fl == read_flags);
	assert(io->fd != STDOUT_FILENO || io->fl == write_flags);

	if (argc > optind) { return usage(program_name); }

	for (i = 0; i < (sizeof io) / (sizeof *io); ++i) {
		/* buffer allocation */
		if ((io[i].buf = malloc(io[i].bs * (sizeof *(io[i].buf)))) == NULL) {
			(void)fprintf(
				stderr, "%s: Failed to allocate %zd bytes\n",
				program_name, io[i].bs
			);
			return EX_OSERR;
		}

		/* easy seeking */
		if (!fdisstd(io[i].fd) && lseek(io[i].fd, io[i].seek, SEEK_SET) != -1) {
			io[i].seek = 0;
		}
	}

	assert(io[1].bufuse == 0); /* requirement for hard seeking */

	/* hard seeking; t is io[1].bufuse, before Io_write subtracts from it */
	for(size_t t; io[1].seek > 0; io[1].seek -= (t - io[1].bufuse)) {
		(void)memset(
			io[1].buf, '\0', /* set buf to all nulls */
			(t = io[1].bufuse = MIN(io[1].bs, io[1].seek)) /* saturate block */
		);

		if (Io_write(&io[1])->bufuse == t && !retry && io[1].error == 0) {
			(void)Io_write(&io[1]); /* second chance */
		}

		if (io[1].error != 0) { return oserr(io[1].fn, io[1].error); }

		if (io[1].bufuse == t) { break; } /* all writes failed! */
	}

	io[1].bufuse = 0; /* reset after hard seek */

	if (io[1].seek > 0) { /* hard seeking failed */
		(void)fprintio(stderr, fmt, io);
		return oserr(io[1].fn, errno);
	}

	for (	;
			count == -1 || count > 0;
			count -= (count != -1) /* decrement if counting */
	) {
		assert(io[0].bufuse == 0);

		{ /* read */
			long skipping;
			size_t t;

			/* hack to intentionally get a partial read from Io_read */
			if ((skipping = MIN(io[0].seek, io[0].bs)) > 0) {
				io[0].bufuse = io[0].bs - (size_t)skipping;
			}

			t = io[0].bufuse;
			if (Io_read(&io[0])->bufuse == t && !retry && io[0].error == 0) {
				(void)Io_read(&io[0]); /* second chance */
			}

			assert(io[0].bufuse >= t);

			if (io[0].bufuse == t) { break; } /* that's all she wrote */

			if (/* t < io[0].bufuse && */ io[0].bufuse < io[0].bs) {
				(void)fprintf(stderr, "%s: Partial read:\n\t", program_name);
				(void)fprintio(stderr, fmt, io);

				if (!retry) { count = 1; }

				if (align >= 0) {
					/* fill the rest of the ibuf with padding */
					(void)memset(
						&(io[0].buf)[io[0].bufuse],
						align,
						io[0].bs - io[0].bufuse
					);

					io->bufuse = io->bs;
				}
			}

			if (skipping > 0) {
				io[0].seek -= skipping;
				io[0].bufuse = 0;
				count += (count != -1); /* increment if counting */
				continue;
			}
		}

		assert(io[0].bufuse > 0);

		while (io[0].bufuse > 0) { /* write */
			if (io[0].bs <= io[1].bs) {
				int n;

				(void)memcpy( /* saturate obuf */
					io[1].buf, io[0].buf,
				   (io[1].bufuse = (n = MIN(io[0].bufuse, io[1].bs)))
				);

				/* permute the copied units out of ibuf */
				(void)memmove(io[0].buf, &(io[0].buf)[n], (io[0].bufuse -= n));
			} else /* if(io[0].bs > io[1].bs) */ {
				int n;

				/* drain what we can from ibuf */
				(void)memcpy(
					&(io[1].buf)[io[1].bufuse], io[0].buf,
				    (n = MIN(io[0].bufuse, io[1].bs - io[1].bufuse))
				);
				io[1].bufuse += n;

				/* permute out the copied units */
				(void)memmove(io[0].buf, &(io[0].buf)[n], io[0].bs - n);
				io[0].bufuse -= n;

				if(io[0].bs + io[1].bufuse <= io[1].bs && count != 1) {
					continue; /* obuf not saturated - we could write more */
				}
			}

			{ /* writes actually happen, or die */
				size_t t;

				t = io[1].bufuse;
				if (Io_write(&io[1])->bufuse == t
						&& !retry
						&& io[1].error == 0) {
					(void)Io_write(&io[1]); /* second chance */
				}

				assert(io[1].error == 0 || io[1].bufuse == t);
				/* if the Io_writes errored, bufuse wouldn't have changed, and
				 * the error will be reported at the end of the read/write
				 * loop */

				assert(io[1].bufuse <= t);

				if (io[1].bufuse == t) { /* no more love */
					count = 1;
					break;
				}
			}

			if (0 < io[1].bufuse /* && io[1].bufuse < t */) {
				(void)fprintf(stderr, "%s: Partial write:\n\t", program_name);
				(void)fprintio(stderr, fmt, io);

				if(!retry) { count = 1; }
			}
		}
	}

	(void)fprintio(stderr, fmt, io);

	for (i = 0; i < (sizeof io) / (sizeof *io); ++i) {
		if (io[i].error) { return oserr(io[i].fn, io[i].error); }
	}

	return EX_OK;
}