#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <sys/file.h>
#include <sys/mman.h>
#include <sys/uio.h>
#include <sys/wait.h>
#include <errno.h>
#include <malloc.h>

#define	BUFSIZE	4096
#define PATH_MAX 4096
#define	MINIOSIZE 4*1024
#define	MAXIOSIZE 4*1024*1024

#ifndef O_DIRECT
#define O_DIRECT 040000
#endif

static int	bufsize = 0;		/* Buffersize. Default random */
static int	bufsize_random = 1;	/* Buffersize random flag. Default random */
static long long	filesize = 0;	/* Target file size */
static long long	target_range = 0;	/* Target range per process */
static int	files_num = 0;
static char	*filenames;
static int	mmap_write = 0;
static int	o_sync = 0;
static int	o_async = 0;
static int	o_direct = 0;

static void setup(void);
static void cleanup(void);

void prg_usage()
{
	fprintf(stderr, "Usage: fs_write [-S 0|1] [-D] [-n num_process]\n");
	fprintf(stderr, "                -s file_size file...\n");
	fprintf(stderr, "\n");
	fprintf(stderr, "		-S	sync mode (0:O_ASYNC 1:O_SYNC)\n");
	fprintf(stderr, "		-D	Direct IO\n");
	fprintf(stderr, "		-n	Process count (default is 1)\n");
	fprintf(stderr, "		-s	File size\n");
	fprintf(stderr, "		file...	Traget file\n");

	exit(1);
}

int runtest(int fd_w[], int childnum)
{
	int num = 0;
	int err = 0;
	long long offset = target_range * childnum;
	long long  remain = target_range;
	int len = bufsize;
	int write_len = 0;
	int psize = getpagesize();
	int max_page = MAXIOSIZE / psize;
	char *buf_w;
	void *p;

	buf_w = (char*)memalign(psize, MAXIOSIZE);

	while (remain > 0) {
		if (bufsize_random == 1) {
			len = (rand() % max_page) * psize;
			if (!len)
				len = MINIOSIZE;
		}
		if (len > remain)
			len = remain;

		for (num = 0; num < files_num ; num++) {
			memset(buf_w, childnum+num, len);
			if (!mmap_write) {
				if (lseek(fd_w[num], offset, SEEK_SET) < 0) {
					err = errno;
					goto out;
				}
				if ((write_len = write(fd_w[num], buf_w, len)) < 0) {
					err = errno;
					goto out;
				}
				if (len != write_len)
					goto out;
			} else {
				p = mmap(NULL, len, PROT_WRITE, MAP_SHARED,
							fd_w[num], offset);
				if (p == MAP_FAILED) {
					err = errno;
					goto out;
				}
				memcpy(p, buf_w, len);
				if (munmap(p, len) == -1) {
					err = errno;
					goto out;
				}
			}
		}
		offset += len;
		remain -= len;
	}

out:
	if (err == ENOSPC || len != write_len)
		err = 0;
	if (err > 0)
		fprintf(stderr, "%s;\n", strerror(err));
	free(buf_w);

	return err;
}

int child_function(int childnum)
{
	int num, ret = -1;
	int fd_w[files_num];
	char *filename;
	
	for (num = 0; num < files_num ; num++) {
		fd_w[num] = -1;
	}
	for (num = 0; num < files_num ; num++) {
		filename = filenames + (num*PATH_MAX);
		if ((fd_w[num] = open(filename,
				O_RDWR|o_sync|o_async|o_direct)) < 0) {
			printf("TINFO:fd_w open failed for %s: %s\n",filename, strerror(errno));
			goto out;
		}
	}

	ret = runtest(fd_w, childnum);

out:
	for (num = 0; num < files_num ; num++) {
		if (fd_w[num] != -1)
			close(fd_w[num]);
	}

	exit(ret);
}

static void setup(void)
{
	int num, fd;
	char *filename;

	for (num = 0; num < files_num; num++) {
		filename = filenames + (num * PATH_MAX);
		if ((fd = open(filename, O_CREAT|O_EXCL|O_WRONLY|O_TRUNC, 0666)) < 0) {
			printf("TBROK, Couldn't create test file %s: %s\n", filename, strerror(errno));
			cleanup();
		}

		if (mmap_write) {
			if (ftruncate(fd, filesize) < 0) {
				printf("TBROK2: Couldn't create test file %s: %s\n", filename, strerror(errno));
				cleanup();
			}
		}
		close(fd);
	}
}

static void cleanup(void)
{
	int num;
	for (num = 0; num < files_num; num++)
		unlink(filenames + (num * PATH_MAX));
	free(filenames);
	exit(1);
}

int main(int argc, char *argv[])
{
	int	numchild = 1;	/* Number of children. Default 5 */
	int	i, num;
	int	err = 0;
	int	sync_mode = 0;
	int	*cpid;
	int 	status;

	while ((i = getopt(argc, argv, "MS:Dn:s:")) != -1) {
		switch(i) {
		case 'M':
			mmap_write = 1;
			break;
		case 'S':
			if ((sync_mode = atoi(optarg)) < 0) {
				fprintf(stderr, "sync mode must be 0 or 1;\n");
				prg_usage();
			}
			switch (sync_mode) {
			case 0:
				o_async = O_ASYNC;
				break;
			case 1:
				o_sync = O_SYNC;
				break;
			default:
				fprintf(stderr, "sync mode must be 0 or 1;\n");
				prg_usage();
			}
			break;
		case 'D':
			o_direct = O_DIRECT;
			break;
		case 'n':
			if ((numchild = atoi(optarg)) <= 0) {
				fprintf(stderr, "no of children must be > 0;\n");
				prg_usage();
			}
			break;
		case 's':
			filesize = atol(optarg);
			if (filesize <= 0) {
				fprintf(stderr, "filesize must be > 0;\n");
				prg_usage();
			}
			break;
		default:
			prg_usage();
		}
	}

	files_num = argc - optind;
	if (files_num <= 0)
		prg_usage();
	if (filesize % numchild != 0 || (filesize / numchild) % BUFSIZE != 0) {
		fprintf(stderr, "filesize must be multiple of 4k*numchild:"
				"filesize=%lld;\n", filesize);
		prg_usage();
	}

	if ((filenames = (char *)malloc(files_num * PATH_MAX)) == NULL ||
	    (cpid = (int *)malloc(sizeof(int)*numchild)) == NULL) {
		perror("malloc");
		exit(1);
	}
	for(i = optind, num = 0 ; i < argc; i++, num++)
		strcpy((char *)filenames + (num * PATH_MAX), argv[i]);
	target_range = filesize / numchild;


	setup();

	for(i = 0; i < numchild; i++) {
		*(cpid+i) = fork();
		if (*(cpid+i) == 0)
			child_function(i);
	}
	for(i = 0; i < numchild; i++) {
		waitpid(*(cpid+1), &status, 0);
	}

	cleanup();
	return err;
}