/* See LICENSE file for copyright and license details. */
#include "common.h"
#include <libsimple-arg.h>


USAGE("[-b blocksize] [-o offset] [-l length | -e postend] [-r] device [< random-source]");
/* TODO document (also in README and man pages) options -b and -r */


struct status status = STATUS_INIT;
_Atomic volatile sig_atomic_t exiting = 0;

static struct span *spans = NULL;
static size_t nspans = 0;
static size_t spans_size = 0;

static off_t total_size = 0;
static char total_size_1000[256];
static char total_size_1024[256];

static clockid_t clck = CLOCK_MONOTONIC_COARSE;
static const char *clkcstr = "CLOCK_MONOTONIC_COARSE";
static struct timespec max_success = {-1, 0};
static struct timespec start_time;

static const struct timespec status_print_interval = {0, MILLISECONDS(500)};
static const struct timespec poll_timeout          = {1, 0};
static int status_print_pipe[2];


static void
add_span(off_t off, off_t amount, size_t blocksize, int try_join)
{
	off_t end = off + amount;

	while ((off_t)(blocksize >> 1) >= amount)
		blocksize >>= 1;

	if (try_join) {
		if (off == spans[nspans - 1U].end || end == spans[nspans - 1U].start) {
			spans[nspans - 1U].start = MIN(off, spans[nspans - 1U].start);
			spans[nspans - 1U].end = MAX(end, spans[nspans - 1U].end);
			spans[nspans - 1U].bad += amount;
			spans[nspans - 1U].blocksize = MAX(blocksize, spans[nspans - 1U].blocksize);
			return;
		}
	}

	if (nspans == spans_size) {
		spans_size += 1024;
		spans = ereallocarray(spans, spans_size, sizeof(*spans));
	}

	spans[nspans].start = off;
	spans[nspans].end = end;
	spans[nspans].bad = amount;
	spans[nspans].blocksize = blocksize;
	nspans++;
}


static void
print_status(int done, struct status *s)
{
	static char buf1[2048] = {0};
	static char buf2[2048] = {0};
	static int bufi = 0;

	char subbuf1[256];
	char subbuf2[256];
	char subbuf3[256];
	char subbuf4[256];
	char subbuf5[256];
	char subbuf6[256];
	char subbuf7[256];
	struct timespec since_success = {-1, 0};
	struct timespec time_spent;

	if (s->last_success.tv_sec >= 0) {
		libsimple_difftimespec(&since_success, &s->now, &s->last_success);
		if (libsimple_cmptimespec(&since_success, &max_success) > 0)
			max_success = since_success;
	}

	/* TODO deal with machine and process suspension */
	libsimple_difftimespec(&time_spent, &s->now, &start_time);

	sprintf(bufi == 0 ? buf1 : buf2,
	        "%ji bytes (%s, %s, %.2lf %%) of %s (%s) shredded\033[K\n"
	        "failed writes: %ju; bad sections: %ju (%ji bytes, %s, %s)\033[K\n"
	        "time spent shredding: %s; performance: %s\033[K\n"
	        "time since last successful write: %s\033[K\n"
	        "maximum time until a successful write: %s\033[K\n"
	        "pass: %ju; pass direction: %s\033[K\n"
	        "%s",
	        /* line 1 { */
	        (intmax_t)s->shredded,
	        humansize1000(s->shredded, subbuf1),
	        humansize1024(s->shredded, subbuf2),
	        100 * (double)s->shredded / (double)total_size,
	        total_size_1000,
	        total_size_1024,
	        /* } line 2 { */
	        s->bad_writes,
	        s->bad_sections,
	        (intmax_t)s->bad_bytes,
	        humansize1000(s->bad_bytes, subbuf3),
	        humansize1024(s->bad_bytes, subbuf4),
	        /* } line 3 { */
	        durationstr(&time_spent, subbuf5, 1),
	        wravg_get(s),
	        /* } line 4 { */
	        durationstr(&since_success, subbuf6, 2),
	        /* } line 5 { */
	        durationstr(&max_success, subbuf7, 2),
	        /* } line 6 { */
	        s->pass_nr,
	        s->direction == FORWARDS ? "forwards" : "backwards",
	        /* } */
	        done ? "" : "\033[6A");

	if (strcmp(buf1, buf2)) {
		fprintf(stderr, "%s", bufi == 0 ? buf1 : buf2);
		fflush(stderr);
	}

	bufi ^= 1;
}


static void *
status_print_loop(void *initial_status)
{
	struct pollfd pfd = {.fd = status_print_pipe[0], .events = POLLIN};
	ssize_t r;
	int terminate = 0;
	struct status s = *(struct status *)initial_status;
	size_t status_off;

	unblock_sigs();

	do {
		switch (ppoll(&pfd, 1U, &poll_timeout, NULL)) {
		case -1:
			if (errno != EINTR)
				eprintf("ppoll:");
			if (exiting)
				break;
			continue;
		case 0:
			break;
		default:
			status_off = 0;
		read_status_again:
			r = read(status_print_pipe[0], &((char *)&s)[status_off], sizeof(s) - status_off);
			if (r == (ssize_t)(sizeof(s) - status_off)) {
				goto have_time;
			} else if (!r) {
				terminate = 1;
				break;
			} else if (r < 0) {
				if (errno == EINTR)
					goto read_status_again;
				eprintf("read <internal pipe>:");
			} else {
				status_off += (size_t)r;
				goto read_status_again;
			}
		}
		if (clock_gettime(clck, &s.now))
			eprintf("clock_gettime %s:", clkcstr);
	have_time:
		if (exiting) {
			fprintf(stderr, "\033[K\nTermination initiated by user...\033[K\n\033[K\n");
			terminate = 1;
		}
		print_status(0, &s);
		wravg_update_time(&s);
	} while (!terminate);

	return NULL;
}


static void
shredspan(int fd, struct span *span, const char *fname)
{
	off_t off, n;
	ssize_t r;
	size_t status_off;
	struct timespec when = {0, 0};
	int bad = span->bad > 0;
	int first_fail = 1;
	int was_successful = 0;
	const char *random_data;

	off = (status.direction == FORWARDS ? span->start : span->end);
	while (status.direction == FORWARDS ? off < span->end : off > span->start) {
		if (exiting) {
		userexit:
			if (status.direction == FORWARDS)
				span->start = off;
			else
				span->end = off;
			close(status_print_pipe[1]);
			status_print_pipe[1] = -1;
			return;
		}

		if (clock_gettime(clck, &status.now))
			eprintf("clock_gettime %s:", clkcstr);
		if (was_successful) {
			was_successful = 0;
			status.last_success = status.now;
		}
		if (libsimple_cmptimespec(&status.now, &when) >= 0) {
			libsimple_sumtimespec(&when, &status.now, &status_print_interval);
			status_off = 0;
		write_status_again:
			r = write(status_print_pipe[1], &((const char *)&status)[status_off], sizeof(status) - status_off);
			if (r != (ssize_t)(sizeof(status) - status_off)) {
				if (r >= 0) {
					status_off += (size_t)r;
					goto write_status_again;
				} else if (r == EINTR) {
					if (exiting)
						goto userexit;
					goto write_status_again;
				} else {
					eprintf("write <internal pipe>:");
				}
			}
		}
		if (status.direction == FORWARDS) {
			n = MIN((off_t)span->blocksize, span->end - off);
		} else {
			n = off--;
			off -= off % (off_t)span->blocksize;
			if (off < span->start)
				off = span->start;
			n -= off;
			if (!n)
				break;
		}
		random_data = get_random(span->blocksize);
	pwrite_again:
		r = pwrite(fd, random_data, (size_t)n, off);
		if (r < 0) {
			if (errno == EINTR) {
				if (exiting)
					goto userexit;
				goto pwrite_again;
			}
			if (errno != EIO)
				eprintf("pwrite %s <buffer> %zu %ji:", fname, (size_t)n, (intmax_t)off);
			add_span(off, n, span->blocksize == 1U ? 1U : span->blocksize >> 1, !first_fail);
			first_fail = 0;
			if (status.direction == FORWARDS)
				off += n;
			if (!span->bad)
				status.bad_bytes += n;
			if (bad)
				bad = 0;
			else
				status.bad_sections += 1U;
			when.tv_sec = 0;
			when.tv_nsec = 0;
			status.bad_writes += 1U;
			continue;
		}
		if (status.direction == FORWARDS) {
			off += (off_t)r;
		} else if ((off_t)r < n) {
			n -= (off_t)r;
			add_span(off + (off_t)r, n, span->blocksize == 1U ? 1U : span->blocksize >> 1, !first_fail);
			first_fail = 0;
			if (status.direction == FORWARDS)
				off += n;
			if (!span->bad)
				status.bad_bytes += n;
			if (bad)
				bad = 0;
			else
				status.bad_sections += 1U;
			when.tv_sec = 0;
			when.tv_nsec = 0;
			status.bad_writes += 1U;
		}
		status.shredded += (off_t)r;
		used_random(r);
		wravg_update_write(r, &status);
		was_successful = 1;
		if (span->bad) {
			status.bad_bytes -= (off_t)r;
			span->bad -= (off_t)r;
		}
	}

	if (bad && !span->bad)
		status.bad_sections -= 1U;
}


static void
dump_map(int fd, const char *fname)
{
	size_t i;
	int r;

	if (!nspans)
		return;

	for (i = 0; i < nspans; i++) {
		r = dprintf(fd, "%s%jx-%jx/%zx",
		            i ? "," : "0x",
		            (uintmax_t)spans[i].start,
		            (uintmax_t)spans[i].end,
		            spans[i].blocksize);
		if (r < 0)
			goto fail;
	}
	r = dprintf(fd, "\n");
	if (r < 0)
		goto fail;

	return;

fail:
	eprintf("dprintf %s:", fname);
}


int
main(int argc, char *argv[])
{
	off_t off = -1, len = -1, end = -1, blk = -1;
	size_t i, j;
	int fd, reverse_direction = 0;
	pthread_t status_print_thread;
	struct status initial_status;

	ARGBEGIN {
	case 'b':
		if (blk >= 0)
			usage();
		blk = unhumansize(ARG(), FLAG());
		break;
	case 'e':
		if (len >= 0 || end >= 0)
			usage();
		end = unhumansize(ARG(), FLAG());
		break;
	case 'l':
		if (len >= 0 || end >= 0)
			usage();
		len = unhumansize(ARG(), FLAG());
		break;
	case 'o':
		if (off >= 0)
			usage();
		off = unhumansize(ARG(), FLAG());
		break;
	case 'r':
		reverse_direction = 1;
		break;
	default:
		usage();
	} ARGEND;

	if (argc != 1)
		usage();

	/* TODO If in foreground, and /dev/tty can be opened,
	 *      use /dev/tty — however if stdin is a tty, use
	 *      stdin instead — to request configuration unless
	 *      bypassed with new option (-Y). The prompt shall
	 *      be identify the device and display it's size,
	 *      also the section the wipe shall be displayed
	 *      if -elo was used. The use shall be required
	 *      to enter 'yes' in upper case. */

	fd = open(argv[0], O_WRONLY | O_DSYNC);
	if (fd < 0)
		eprintf("open %s O_WRONLY|O_DSYNC:", argv[0]);

	init_random(STDIN_FILENO, "<stdin>", blk);

	spans = emalloc(sizeof(*spans));
	spans[0].start = 0;
	spans[0].end = filesize(fd, argv[0]);
	spans[0].bad = 0;
	spans[0].blocksize = max_blksize();
	nspans = 1U;
	spans_size = 1U;

	if (off >= 0) {
		if (off > spans[0].end)
			eprintf("value of -o flag is beyond the end of the file");
		spans[0].start = off;
	}
	if (len >= 0) {
		if (len > OFF_MAX - spans[0].start)
			eprintf("the sum of the values of -o and -l flag is too large");
		end = spans[0].start + len;
		if (end > spans[0].end)
			eprintf("the sum of the values of -o and -l flag is beyond the end of the file");
		spans[0].end = end;
	} else if (end >= 0) {
		if (end > spans[0].end)
			eprintf("the value of -e flag is beyond the end of the file");
		spans[0].end = end;
	}

	total_size = spans[0].end - spans[0].start;
	if (!total_size) {
		weprintf("attempting to shred 0 bytes");
		close(fd);
		return 0;
	}

	humansize1000(total_size, total_size_1000);
	humansize1024(total_size, total_size_1024);

	if (pipe(status_print_pipe))
		eprintf("pipe:");

	setup_sighandler();
	block_sigs();

	if (reverse_direction)
		status.direction ^= 1;
	if (clock_gettime(clck, &start_time)) {
		clck = CLOCK_MONOTONIC;
		clkcstr = "CLOCK_MONOTONIC";
		if (clock_gettime(clck, &start_time))
			eprintf("clock_gettime %s:", clkcstr);
	}
	wravg_init(&start_time, &status);
	status.last_success = start_time;

	initial_status = status;
	errno = pthread_create(&status_print_thread, NULL, &status_print_loop, &initial_status);
	if (errno)
		eprintf("pthread_create NULL:");

	for (;;) {
		size_t old_nspans = nspans;
		for (i = 0; i < old_nspans && !exiting; i++)
			shredspan(fd, &spans[i], argv[0]);
		if (exiting) {
			memmove(&spans[0], &spans[i], (nspans -= i) * sizeof(*spans));
			break;
		}
		for (i = 0, j = nspans, nspans -= old_nspans; i < nspans;)
			spans[i++] = spans[--j];
		if (!nspans)
			break;
		status.direction ^= 1;
		status.pass_nr += 1U;
	}

	close(fd);
	if (status_print_pipe[1] >= 0)
		close(status_print_pipe[1]);
	destroy_random();

	errno = pthread_join(status_print_thread, NULL);
	if (errno)
		weprintf("pthread_join:");

	if (clock_gettime(clck, &status.now))
		eprintf("clock_gettime %s:", clkcstr);
	print_status(1, &status);

	if (nspans) {
		/* TODO document in man page */
		dump_map(STDOUT_FILENO, "<stdout>");
		if (close(STDOUT_FILENO))
			eprintf("write <stdout>");
	}
	return 0;
}