/*
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License version 2 as
 * published by the Free Software Foundation.
 *
 * Copyright 2010, Coly Li <i@coly.li>
 */


#define _GNU_SOURCE
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/time.h>
#include <linux/fs.h>
#include <stdint.h>
#include <fcntl.h>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <dirent.h>
#include <malloc.h>
#include <stropts.h>

void usage()
{
	printf("Usage: seekrw -f <file path> -a <access times> -l <read/write length> -r -w -d\n"
	       "       -r read\n"
	       "       -w write\n"
	       "       -d directio\n");
	exit(1);
}

int check_io_size(int io_size)
{
	int bits = 0;
	printf("%s:%d io_size: %d\n", __func__, __LINE__, io_size);
	while ((io_size = io_size >> 1)) {
		bits ++;
	
	}

	io_size = 1 << bits;
	printf("%s:%d io_size: %d\n", __func__, __LINE__, io_size);
	return io_size;
}

size_t get_blk_fsize(char *path, int io_size)
{
	int fd;
	size_t size;

	fd = open64(path, O_RDONLY);
	if (fd < 0) {
		printf("open block device file %s failed.\n", path);
		exit(1);
	}

	if (ioctl(fd, BLKGETSIZE64, &size) < 0) {
		printf("ioctl BLKGETSIZE64 to block device file %s failed.\n", path);
		exit(1);
	}

	return size;
}

size_t check_file(char *path, int io_size)
{
	int ret;
	struct stat statbuf;
	unsigned long limit = 1;

	size_t f_size;

	ret = stat(path, &statbuf);
	if (ret) {
		printf("stat file %s failed (%s).\n", path, strerror(errno));
		exit(1);
	}

	if (S_ISREG(statbuf.st_mode)) {
		f_size = statbuf.st_size;
		printf("f_size: %lu, limit: %lu\n", (unsigned long)f_size, (unsigned long)(limit *(1<<30)));
		if (f_size < (limit * (1<<30)) || f_size <= io_size) {
			printf("file is should be larger than %lu GB, and be larger than read/write length %d bytes\n", limit, io_size);
			exit(1);
		}
	} else if (S_ISBLK(statbuf.st_mode)) {
		printf("check block device file %s\n", path);
		f_size = get_blk_fsize(path, io_size);
	} else {
		printf("unsupported file type.\n");
		exit(1);
	}

	f_size = (f_size / io_size) * io_size;
	return f_size;
}

void prepare_random_offset(off_t * random_buf, int access_loop, size_t f_size, int io_size)
{
	int i;
	uint32_t off1, off2;
	uint64_t off;

	for (i = 0; i < access_loop; i ++) {
		off1 = (uint32_t)random();
		off2 = (uint32_t)random();
		off = (uint64_t)(((uint64_t)off1 << 32) + off2);
		off = (uint64_t)(off % f_size);
		off = (off / io_size) * io_size;
		random_buf[i] = (off_t)off;
	}
}

void do_seek_read_write(int fd, uint64_t offset, char *io_buf, int io_size, int rw)
{
	off_t ret;

	printf("%s %d bytes %s offset 0x%016llx.\n",
		rw ? "write" : "read", io_size,
	       	rw ? "to" : "from",
		(unsigned long long)offset);

	if (rw)
		/* write */
		ret = pwrite64(fd, io_buf, io_size, offset);
	else
		/* read */
		ret = pread64(fd, io_buf, io_size, offset);

	if (ret < 0) {
		printf("%s failed on offset %llu (%s)\n",
			rw ? "write":"read", (unsigned long long)offset, strerror(errno));
		exit(1);
	}
}

void time_diff(struct timeval tv1, struct timeval tv2)
{
	unsigned long diff_m, diff_s, diff_us;

	if (tv2.tv_usec < tv1.tv_usec) {
		tv2.tv_usec += 1000000;
		tv2.tv_sec --;
	}

	diff_us = tv2.tv_usec - tv1.tv_usec;
	diff_s = tv2.tv_sec - tv1.tv_sec;
	diff_m = diff_s / 60;
	diff_s = diff_s % 60;

	printf("real time:");
	if (diff_m)
		printf(" %lu minutes", diff_m);
	if (diff_s)
		printf(" %lu seconds", diff_s);
	if (diff_us)
		printf(" %lu usec", diff_us);
	printf("\n");
}


int main(int argc, char *argv[])
{
	char path[4096] = {0,};
	size_t f_size;
	int io_size = 0;
	int io_limit = 1;
	int rw = -1, dio = 0;
	off_t *random_offset_buf;
	char *io_buf;
	unsigned long max_access_loop = 0xfffffffffffffff0;
	long int access_loop;
	struct timeval t1, t2;
	int fd, i;
	char c;

	while((c = getopt(argc, argv, "hf:l:a:rwd")) != EOF)
	{
		switch(c)
		{
			case 'h':
				usage();
				break;
			case 'f':
				strncpy(path, optarg, sizeof(path));
				break;
			case 'l':
				io_size = atoi(optarg);
				if (io_size > io_limit*(1<<20)) {
					printf("read/write length should be less then %d MB\n", io_limit);
					exit(1);
				}
				if (io_size < 512) {
					printf("read/write length should at least be 512 B\n)");
					exit(1);
				}
				if (io_size <=0)
					usage();
				break;
			case 'a':
				access_loop = atoi(optarg);
				if (access_loop > max_access_loop) {
					printf("access loop number should be less than %lu\n", max_access_loop);
					exit(1);
				}
				if (access_loop <= 0)
					usage();
				break;
			case 'w':
				rw ++;
			case 'r':
				rw ++;
				break;
			case 'd':
				dio = 1;
				break;
			default:
				usage();
				break;
		}
	}

	if (strlen(path) == 0 || io_size == 0 || rw < 0 || rw > 1)
		usage();

	/* add file name */
	printf("file path %s\n", path);

	io_size = check_io_size(io_size);
	f_size = check_file(path, io_size);

	io_buf = (char *)memalign(1<<20, io_size);
	if (io_buf == NULL) {
		printf("malloc to (%d bytes) io_buf failed (%s)\n", io_size, strerror(errno));
		exit(1);
	}

	random_offset_buf = (off_t *)malloc(access_loop * sizeof(off_t));
	if (random_offset_buf == NULL) {
		printf("malloc to random offset buffer failed (%s)\n", strerror(errno));
		exit(1);
	}

	prepare_random_offset(random_offset_buf, access_loop, f_size, io_size);

	if (dio)
		fd = open(path, O_RDWR | O_DIRECT);
	else
		fd = open(path, O_RDWR);
	if (fd < 0) {
		printf("open file %s failed (%s).\n", path, strerror(errno));
		exit(1);
	}

	gettimeofday(&t1, NULL);
	for(i = 0; i < access_loop; i++) {
		do_seek_read_write(fd, random_offset_buf[i], io_buf, io_size, rw);
	}
	fsync(fd);
	gettimeofday(&t2, NULL);

	close(fd);

	free(random_offset_buf);
	free(io_buf);

	time_diff(t1, t2);
	return 0;
}
