/*
 * Copyright 2009, Ingo Weinhold, ingo_weinhold@gmx.de.
 * Copyright 2007, Marcus Overhagen. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 */

#include <errno.h>
#include <fcntl.h>
#include <getopt.h>
#include <limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>


#define EXIT_FAILURE 1


static void
print_usage(bool error)
{
	printf("\n");
	printf("create_image\n");
	printf("\n");
	printf("usage: create_image -i <imagesize> [-c] [-f] <file>\n");
	printf("       -i, --imagesize    size of raw partition image file\n");
	printf("       -f, --file         the raw partition image file\n");
	printf("       -c, --clear-image  set the image content to zero\n");
	exit(error ? EXIT_FAILURE : 0);
}


int
main(int argc, char *argv[])
{
	off_t imageSize = 0;
	const char *file = NULL;
	bool clearImage = false;

	while (1) {
		int c;
		static struct option long_options[] = {
			{"file", required_argument, 0, 'f'},
			{"clear-image", no_argument, 0, 'c'},
			{"help", no_argument, 0, 'h'},
			{"imagesize", required_argument, 0, 'i'},
			{0, 0, 0, 0}
		};

		opterr = 0; /* don't print errors */
		c = getopt_long(argc, argv, "+hi:cf:", long_options, NULL);
		if (c == -1)
			break;

		switch (c) {
			case 'h':
				print_usage(false);
				break;

			case 'i':
				imageSize = strtoull(optarg, NULL, 10);
				if (strchr(optarg, 'G') || strchr(optarg, 'g'))
					imageSize *= 1024 * 1024 * 1024;
				else if (strchr(optarg, 'M') || strchr(optarg, 'm'))
					imageSize *= 1024 * 1024;
				else if (strchr(optarg, 'K') || strchr(optarg, 'k'))
					imageSize *= 1024;
				break;

			case 'f':
				file = optarg;
				break;

			case 'c':
				clearImage = true;
				break;

			default:
				print_usage(true);
		}
	}

	if (file == NULL && optind == argc - 1)
		file = argv[optind];

	if (!imageSize || !file)
		print_usage(true);

	if (imageSize < 0) {
		fprintf(stderr, "Error: invalid image size\n");
		exit(EXIT_FAILURE);
	}

	if (imageSize % 512) {
		fprintf(stderr, "Error: image size must be a multiple of 512 bytes\n");
		exit(EXIT_FAILURE);
	}

	int fd = open(file, O_RDWR | O_CREAT, 0666);
	if (fd < 0) {
		fprintf(stderr, "Error: couldn't open file %s (%s)\n", file,
			strerror(errno));
		exit(EXIT_FAILURE);
	}

	struct stat st;
	if (fstat(fd, &st) < 0) {
		fprintf(stderr, "Error: stat()ing file %s failed (%s)\n", file,
			strerror(errno));
		exit(EXIT_FAILURE);
	}

	if (!S_ISREG(st.st_mode) && !S_ISBLK(st.st_mode) && !S_ISCHR(st.st_mode)) {
		fprintf(stderr, "Error: type of file %s not supported\n", file);
		exit(EXIT_FAILURE);
	}

	if (S_ISREG(st.st_mode)) {
		// regular file -- use ftruncate() to resize it
		if ((clearImage && ftruncate(fd, 0) != 0)
			|| ftruncate(fd, imageSize) != 0) {
			fprintf(stderr, "Error: resizing file %s failed (%s)\n", file,
				strerror(errno));
			exit(EXIT_FAILURE);
		}
	} else {
		// some kind of device -- clear it manually, if we have to
		if (clearImage) {
			char buffer[1024 * 1024];
			memset(buffer, 0, sizeof(buffer));

			off_t totalWritten = 0;
			ssize_t written;
			while ((written = write(fd, buffer, sizeof(buffer))) > 0)
				totalWritten += written;

			// Only fail, if an error occurs and we haven't written anything at
			// all yet.
			// TODO: We should probably first determine the size of the device
			// and try to write only that much.
			if (totalWritten == 0 && written < 0) {
				fprintf(stderr, "Error: writing to device file %s failed "
					"(%s)\n", file, strerror(errno));
				exit(EXIT_FAILURE);
			}
		}
	}

	close(fd);
	return 0;
}