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


USAGE("device [< random-source]");


struct span {
	off_t start;
	off_t end;
	off_t bad;
	size_t blocksize;
};


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

static off_t shredded = 0;
static off_t total_size = 0;

static char reservoir[128U << 10];
static size_t reservoir_off = sizeof(reservoir);
static int use_stdin;

static char total_size_1000[256];
static char total_size_1024[256];

static uintmax_t bad_writes = 0;
static uintmax_t bad_sections = 0;
static off_t bad_bytes = 0;


static off_t
filesize(int fd, const char *fname)
{
	struct stat st;

	if (fstat(fd, &st))
		eprintf("fstat %s:", fname);

	switch (st.st_mode & S_IFMT) {
	case S_IFREG:
		break;
	case S_IFBLK:
		if (ioctl(fd, BLKGETSIZE64, &st.st_size) < 0)
			eprintf("ioctl %s BLKGETSIZE64:", fname);
		break;
	default:
		eprintf("%s: not a regular file or block device", fname);
	}

	return st.st_size;
}


static void
ensure_random(size_t needed)
{
	size_t off;
	ssize_t r;

	if (sizeof(reservoir) - reservoir_off >= needed)
		return;

	if (!use_stdin) {
		libsimple_random_bytes(&libsimple_random_bits, NULL, reservoir, reservoir_off);
		reservoir_off = 0;
		return;
	}

	for (off = 0; off < reservoir_off;) {
		r = read(STDIN_FILENO, &reservoir[off], reservoir_off - off);
		if (r <= 0) {
			if (!r)
				eprintf("random source depleted");
			if (errno == EINTR)
				continue;
			eprintf("read <stdin>:");
		}
		off += (size_t)r;
	}

	reservoir_off = 0;
}


static off_t
add_span(off_t off, off_t max_end, size_t blocksize)
{
	off_t end;

	end = MIN((off | ((off_t)blocksize - 1)) + 1, max_end);
	while ((off_t)(blocksize >> 1) >= end - off)
		blocksize >>= 1;

	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 = end - off;
	spans[nspans].blocksize = (size_t)blocksize;
	nspans++;

	return end - off;
}


static char *
humansize1000(off_t s, char *buf)
{
	const char *units = "kMGTPEZYRQ";
	size_t unit = 0;
	if (s < 1000) {
		sprintf(buf, "%u B", (unsigned)s);
		return buf;
	}
	s /= 100;
	while (units[unit + 1U] && s >= 10000) {
		s /= 1000;
		unit++;
	}
	sprintf(buf, "%u.%u %cB", (unsigned)s / 10U, (unsigned)s % 10U, units[unit]);
	return buf;
}


static char *
humansize1024(off_t s, char *buf)
{
	const char *units = "KMGTPEZYRQ";
	size_t unit = 0;
	if (s < 1024) {
		sprintf(buf, "%u B", (unsigned)s);
		return buf;
	}
	while (units[unit + 1U] && s >= 1024 * 1024) {
		s /= 1024;
		unit++;
	}
	sprintf(buf, "%lu.%lu %ciB", (unsigned long int)s / 1024UL, (unsigned long int)(s * 10 % 10240) / 1024UL, units[unit]);
	return buf;
}


static void
print_progress(int done)
{
	static char buf1[1024] = {0};
	static char buf2[1024] = {0};
	static int i = 0;

	char subbuf1[256];
	char subbuf2[256];
	char subbuf3[256];
	char subbuf4[256];

	sprintf(i == 0 ? buf1 : buf2,
	        "%ji bytes (%s, %s, %.2lf %%) of %s (%s) shredded\033[K\n"
		"failed writes: %ju; bad sections: %ju (%s, %s)\033[K\n%s",
		(intmax_t)shredded,
	        humansize1000(shredded, subbuf1),
	        humansize1024(shredded, subbuf2),
	        100 * (double)shredded / (double)total_size,
	        total_size_1000,
	        total_size_1024,
		bad_writes,
		bad_sections,
	        humansize1000(bad_bytes, subbuf3),
	        humansize1024(bad_bytes, subbuf4),
	        done ? "" : "\033[A\033[A");

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

	i ^= 1;
}


static void
shredspan(int fd, struct span *span, const char *fname)
{
	off_t off, n;
	ssize_t r;
	clockid_t clck = CLOCK_MONOTONIC_COARSE;
	const char *clkcstr = "CLOCK_MONOTONIC_COARSE";
	struct timespec now, when = {0, 0};
	int bad = span->bad > 0;

	if (clock_gettime(clck, &now)) {
		clck = CLOCK_MONOTONIC;
		clkcstr = "CLOCK_MONOTONIC";
	}

	for (off = span->start; off < span->end;) {
		if (clock_gettime(clck, &now))
			eprintf("clock_gettime %s:", clkcstr);
		if (now.tv_sec > when.tv_sec || (now.tv_sec == when.tv_sec && now.tv_nsec >= when.tv_nsec)) {
			when = now;
			when.tv_nsec += 500000000L;
			when.tv_sec += when.tv_nsec / 1000000000L;
			when.tv_nsec = when.tv_nsec % 1000000000L;
			print_progress(0);
		}
		if (lseek(fd, off, SEEK_SET) < 0)
			eprintf("lseek %s %ji SEEK_SET\n", fname, (intmax_t)off);
		ensure_random(span->blocksize);
		r = write(fd, &reservoir[reservoir_off], (size_t)MIN((off_t)span->blocksize, span->end - off));
		if (r < 0) {
			if (errno != EINTR) {
				off += n = add_span(off, span->end, span->blocksize == 1U ? 1U : span->blocksize);
				if (!span->bad)
					bad_bytes += n;
				if (bad)
					bad = 0;
				else
					bad_sections += 1U;
			}
			when.tv_sec = 0;
			when.tv_nsec = 0;
			bad_writes += 1U;
			continue;
		}
		off += (off_t)r;
		shredded += (off_t)r;
		reservoir_off += (size_t)r;
		if (span->bad) {
			bad_bytes -= (off_t)r;
			span->bad -= (off_t)r;
		}
	}

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


int
main(int argc, char *argv[])
{
	size_t i;
	int fd;

	ARGBEGIN {
	default:
		usage();
	} ARGEND;

	if (argc != 1)
		usage();

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

	use_stdin = !isatty(STDIN_FILENO);
	if (!use_stdin) {
		libsimple_srand();
	}

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

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

	while (nspans) {
		size_t old_nspans = nspans;
		for (i = 0; i < old_nspans; i++)
			shredspan(fd, &spans[i], argv[0]);
		nspans -= old_nspans;
		memmove(&spans[0], &spans[old_nspans], nspans * sizeof(*spans));
	}

	close(fd);

	print_progress(1);
	return 0;
}