// code by SASANO Takayoshi, CC BY-SA

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <endian.h>
#include <fcntl.h>
#include <inttypes.h>
#include <time.h>
#include <unistd.h>
#include <sys/ioctl.h>

#if defined(__OpenBSD__)
#include <sys/dkio.h>
#include <sys/disklabel.h>
#elif defined(__linux__)
#include <linux/fs.h>
#endif

#define	BUFFERSIZE	(1024 * 1024 * 1)
#define	SECTORSIZE	512
#define	SECTORS		(BUFFERSIZE / SECTORSIZE)
#define	ELEMENTS	(SECTORSIZE / sizeof(uint64_t))

#define	TEST_MASK		0xfc00
#define	TEST_VERBOSE_STATUS	0x0400
#define	TEST_NODATA		0x0800
#define	TEST_WRITE		0x1000
#define	TEST_READ		0x2000
#define	TEST_SET		0x4000
#define	TEST_ENDLESS		0x8000

#define	DATA_MASK		0x0003
#define	DATA_SECNO		0x0000
#define	DATA_FIXVALUE		0x0001
#define	DATA_RANDOM0		0x0002
#define	DATA_RANDOM1		0x0003
#define	set_FIXVALUE(x)		(((x) & 0xffff) << 16)
#define	get_FIXVALUE(x)		(((x) >> 16) & 0xffff)

#define	BITWIDTH_MASK		0x000c
#define	BITWIDTH_64		0x0000
#define	BITWIDTH_32		0x0004
#define	BITWIDTH_16		0x0008
#define	BITWIDTH_8		0x000c

#define	EFFECT_MASK		0x00e0
#define	EFFECT_BITREVERSE	0x0020
#define	EFFECT_NOT		0x0040
#define	EFFECT_BIGENDIAN	0x0080

static	uint8_t	buf[BUFFERSIZE];
static	uint32_t	myrandom_state;
static	int	progress_percentage;

/* show progress */
static	void	disp_progress(uint64_t curr, uint64_t max)
{
	int	percent;
	time_t	t;
	char	tmp[40];

	percent = (double)(curr * 100) / max;

	if (progress_percentage != percent) {
		t = time(NULL);
		ctime_r(&t, tmp);
		tmp[strlen(tmp) - 1] = '\0';
		printf("[%s] %d%% done\n", tmp, percent);
		progress_percentage = percent;
	}
}

/* pseudo random */
static	uint8_t	myrandom(int v1, int v2)
{
	myrandom_state *= v1;
	myrandom_state += v2;
	return myrandom_state >> 24;
}

/* data generator */
static	uint64_t	generator(int method, uint64_t value)
{
	int	i, m;

	m = method & DATA_MASK;

	switch (m) {
	case	DATA_SECNO:
		/* do nothing */
		break;
	case	DATA_FIXVALUE:
		value = get_FIXVALUE(method);
		break;
	case	DATA_RANDOM0:
		for (i = 0; i < 8; i++) {
			value = (value << 8) | myrandom(13, 0x1a2b3c4d);
		}
		break;
	default:
		for (i = 0; i < 4; i++) {
			value = (value << 8) | myrandom(1, 0xff000000);
			value = (value << 8) | ((value & 0xff) ^ 0xff);
		}
		break;
	}

	return value;
}

/* bitwidth */
static	uint64_t	bitwidth(int method, uint64_t value)
{
	switch (method & BITWIDTH_MASK) {
	case	BITWIDTH_8:
		value = (uint8_t)value;
		value |= (value << 8);
		/* FALLTHROUGH */
	case	BITWIDTH_16:
		value = (uint16_t)value;
		value |= (value << 16);
		/* FALLTHROUGH */
	case	BITWIDTH_32:
		value = (uint32_t)value;
		value |= (value << 32);
		/* FALLTHROUGH */
	default:
	case	BITWIDTH_64:
		break;
	}

	return value;
}

/* bit reverse */
static	uint64_t	bit_reverse(uint64_t v)
{
	v =  (((v & 0x00000000ffffffffULL) << 32) |
	      ((v & 0xffffffff00000000ULL) >> 32));
	v =  (((v & 0x0000ffff0000ffffULL) << 16) |
	      ((v & 0xffff0000ffff0000ULL) >> 16));
	v =  (((v & 0x00ff00ff00ff00ffULL) << 8) |
	      ((v & 0xff00ff00ff00ff00ULL) >> 8));
	v =  (((v & 0x0f0f0f0f0f0f0f0fULL) << 4) |
	      ((v & 0xf0f0f0f0f0f0f0f0ULL) >> 4));
	v =  (((v & 0x3333333333333333ULL) << 2) |
	      ((v & 0xccccccccccccccccULL) >> 2));
	v =  (((v & 0x5555555555555555ULL) << 1) |
	      ((v & 0xaaaaaaaaaaaaaaaaULL) >> 1));

	return v;
}

/* create test data from sector value */
static	uint64_t	test_data(uint64_t secno, int method)
{
	uint64_t	value;

	value = generator(method, secno);
	value = bitwidth(method, value);
	value = (method & EFFECT_BIGENDIAN) ? be64toh(value) : le64toh(value);
	if (method & EFFECT_NOT) value = ~value;
	if (method & EFFECT_BITREVERSE) value = bit_reverse(value);

	return value;
}

/* fill sector with test data */
static	void	fill_sector(uint64_t *p, void **q, uint64_t secno, int method)
{
	int	i;

	for (i = 0; i < ELEMENTS; i++) {
		*p++ = test_data(secno, method);
	}

	*q = p;
}

/* fill buffer with test data */
static	void	fill_buffer(uint64_t secno, uint64_t sectors, int method)
{
	int	i;
	void	*p;

	p = buf;
	
	for (i = 0; i < sectors; i++) {
		fill_sector(p, &p, secno, method);
		secno++;
	}
}

/* dump */
static	void	dump(void *d, int size)
{
	int	i;
	uint8_t	*p;

	p = d;

	for (i = 0; i < size; i++) {
		printf(" %02x", *p);
		p++;
	}
}

/* check sector with test data */
static	int	check_sector(uint64_t *p, void **q, uint64_t secno, int method, int verbose)
{
	int	i, result;
	uint64_t	data;

	result = 0;

	for (i = 0; i < ELEMENTS; i++) {
		data = test_data(secno, method);

		if (data != *p) {
			if (verbose) {
				printf("offset %ld:",
				       i * (SECTORSIZE / ELEMENTS));
				dump(&data, sizeof(data));
				printf(" ->");
				dump(p, sizeof(data));
				printf("\n");
			}
			result = 1;
		}
		p++;
	}

	*q = p;

	return result;
}

/* check buffer with test data */
static	void	check_buffer(uint64_t secno, uint64_t sectors, int method)
{
	int	i;
	void	*p, *q;
	uint32_t	saved_state;

	p = buf;

	for (i = 0; i < sectors; i++) {
		saved_state = myrandom_state;
		if (check_sector(p, &q, secno, method, 0)) {
			printf("sector %"PRId64": check failed\n", secno);
			myrandom_state = saved_state;
			check_sector(p, &q, secno, method, 1);
		}
		secno++;
		p = q;
	}
}

/* write test */
static	void	write_test(int fd, int method, uint64_t sectors)
{
	uint64_t	i, remain;

	myrandom_state = 0;
	progress_percentage = -1;

	for (i = 0; i < sectors; i += SECTORS) {
		if (method & TEST_VERBOSE_STATUS) disp_progress(i, sectors);

		remain = sectors - i;
		if (remain > SECTORS) {
			remain = SECTORS;
		}

		if (!(method & TEST_NODATA)) {
			fill_buffer(i, remain, method);
		}

		if (pwrite(fd, buf, remain * SECTORSIZE,
			   i * SECTORSIZE) < remain * SECTORSIZE) {
			printf("sector %"PRId64": write error\n", i);
		}
	}

	sync();
	if (method & TEST_VERBOSE_STATUS) disp_progress(i, sectors);

	return;
}

/* read test */
static	void	read_test(int fd, int method, uint64_t sectors)
{
	uint64_t	i, remain;

	myrandom_state = 0;
	progress_percentage = -1;

	for (i = 0; i < sectors; i += SECTORS) {
		if (method & TEST_VERBOSE_STATUS) disp_progress(i, sectors);

		remain = sectors - i;
		if (remain > SECTORS) {
			remain = SECTORS;
		}

		if (pread(fd, buf, remain * SECTORSIZE,
			  i * SECTORSIZE) < remain * SECTORSIZE) {
			printf("sector %"PRId64": read error\n", i);
		}

		if (!(method & TEST_NODATA)) {
			check_buffer(i, remain, method);
		}
	}

	if (method & TEST_VERBOSE_STATUS) disp_progress(i, sectors);

	return;
}

static	void	sync_sec(void)
{
	time_t	t;

	t = time(NULL);

	while (t == time(NULL));
}

/* get total sectors */
static	uint64_t	get_sectors(int fd)
{
	int	n;
	off_t	size;
	uint64_t	sectors;
#if defined(__OpenBSD__)
	struct disklabel dl;
#elif defined(__linux__)
	unsigned long	v;
#endif
	

#if defined(__OpenBSD__)
	n = ioctl(fd, DIOCGPDINFO, &dl);
#elif defined(__linux__)
	n = ioctl(fd, BLKGETSIZE, &v);
#endif

	/* get total sectors */
	if (n < 0) {
		size = lseek(fd, 0, SEEK_END);
		lseek(fd, 0, SEEK_SET);
		sectors = size / SECTORSIZE;
	} else {
#if defined(__OpenBSD__)
		sectors = DL_GETDSIZE(&dl);
		n = dl.d_secsize % SECTORSIZE;
#elif defined(__linux__)
		sectors = v;
		n = ioctl(fd, BLKSSZGET, &v);
		if (n >= 0) n = v % SECTORSIZE;
#endif
		if (n) {
			printf("device should be %d bytes/sector\n",
			       SECTORSIZE);
			sectors = 0;
			goto fin0;
		}
	}

fin0:
	return sectors;
}

static	void	show_result(char *type, uint64_t sectors, time_t seconds)
{
	uint64_t	bytes;
	double	rate;

	bytes = sectors * SECTORSIZE;

	printf("%s %"PRId64" bytes, %"PRId64" sec", type, bytes, seconds);

	if (seconds) {
		rate = (double)bytes / seconds;
		printf(" (%f bytes/sec)", rate);
	}

	printf("\n");
}

static	void	test(char *device, int method)
{
	int	fd;
	uint64_t	sectors;
	time_t	before, after;

	fd = open(device, O_RDWR);
	if (fd < 0) {
		printf("device open error\n");
		goto fin0;
	}

	sectors = get_sectors(fd);
	printf("device %s: %"PRId64" sectors method %#x\n", device, sectors, method);

	if (method & TEST_WRITE) {
		sync_sec();
		before = time(NULL);
		write_test(fd, method, sectors);
		after = time(NULL);
		show_result("write", sectors, after - before);
	}

	if (method & TEST_READ) {
		sync_sec();
		before = time(NULL);
		read_test(fd, method, sectors);
		after = time(NULL);
		show_result("read", sectors, after - before);
	}

fin0:
	close(fd);
	return;
}

int	main(int argc, char *argv[])
{
	int	method;

	if (argc < 3) {
		printf("usage:	%s [method] [device]\n", argv[0]);
		return 0;
	}

	memset(buf, 0, BUFFERSIZE);
	method = strtol(argv[1], NULL, 0);

	do {
		if (method & TEST_SET) {
			test(argv[2], TEST_WRITE | TEST_READ |
			     DATA_FIXVALUE | BITWIDTH_8 | EFFECT_NOT);
			test(argv[2], TEST_WRITE | TEST_READ | DATA_SECNO);
			test(argv[2], TEST_WRITE | TEST_READ | DATA_RANDOM0);
			test(argv[2], TEST_WRITE | TEST_READ |
			     DATA_FIXVALUE | BITWIDTH_8);
		} else {
			test(argv[2], method);
		}
	} while (method & TEST_ENDLESS);

	return 0;
}