/* 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;
	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 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].blocksize = (size_t)blocksize;
	nspans++;
}


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.%u %ciB", (unsigned long int)s / 1024U, (unsigned long int)((s * 10 % 10240) / 1024U), units[unit]);
	return buf;
}


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

	char subbuf1[256];
	char subbuf2[256];

	sprintf(i == 0 ? buf1 : buf2,
	        "%.2lf %% (%s, %s) of %s (%s) shredded\033[K\n%s",
	        100. * (double)shredded / (double)total_size,
	        humansize1000(shredded, subbuf1),
	        humansize1024(shredded, subbuf2),
	        total_size_1000,
	        total_size_1024,
	        done ? "" : "\033[A");

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

	i ^= 1;
}


static void
shredspan(int fd, const struct span *span, const char *fname)
{
	off_t off;
	ssize_t r;

	for (off = span->start; off < span->end;) {
		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 = add_span(off, span->end, span->blocksize == 1U ? 1U : span->blocksize);
			continue;
		}
		off += (off_t)r;
		shredded += (off_t)r;
		reservoir_off += (size_t)r;
	}
}


int
main(int argc, char *argv[])
{
	char buf1[256];
	char buf2[256];
	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].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;
}