/* Copyright 2011 Fredrik Wikstrom. All rights reserved.
**
** Redistribution and use in source and binary forms, with or without
** modification, are permitted provided that the following conditions
** are met:
**
** 1. Redistributions of source code must retain the above copyright
**    notice, this list of conditions and the following disclaimer.
**
** 2. Redistributions in binary form must reproduce the above copyright
**    notice, this list of conditions and the following disclaimer in the
**    documentation and/or other materials provided with the distribution.
**
** THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS `AS IS'
** AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
** IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
** ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
** LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
** CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
** SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
** INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
** CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
** ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
** POSSIBILITY OF SUCH DAMAGE.
*/

#include "vditool.h"
#include "vditool_rev.h"

CONST USED TEXT verstag[] = VERSTAG;

BPTR StdOut;
BPTR StdErr;

static void show_version (void);
static void show_usage (void);
static int vdi_create (int argc, char **argv);
static int vdi_info (int argc, char **argv);
static int vdi_compact (int argc, char **argv);
static int vdi_convertfromraw (int argc, char **argv);
static int vdi_converttoraw (int argc, char **argv);
static BOOL is_zero_block (const uint8 *buffer, uint32 block_size);

int main (int argc, char **argv) {
	int retval = RETURN_FAIL;
	CONST_STRPTR command;

	StdOut = IDOS->Output();
	StdErr = IDOS->ErrorOutput();
	if (StdErr == ZERO) {
		StdErr = StdOut;
	}

	show_version();

	if (argc < 2 || (argc == 2 && argv[1][0] == '?')) {
		show_usage();
		return RETURN_OK;
	}

	if (argc < 3) {
		show_usage();
		return RETURN_FAIL;
	}

	command = argv[1];
	if (!IUtility->Stricmp(command, "create")) {
		retval = vdi_create(argc, argv);
	} else if (!IUtility->Stricmp(command, "info")) {
		retval = vdi_info(argc, argv);
	} else if (!IUtility->Stricmp(command, "compact")) {
		retval = vdi_compact(argc, argv);
	} else if (!IUtility->Stricmp(command, "convertfromraw")) {
		retval = vdi_convertfromraw(argc, argv);
	} else if (!IUtility->Stricmp(command, "converttoraw")) {
		retval = vdi_converttoraw(argc, argv);
	} else {
		show_usage();
	}

	return retval;
}

static void show_version (void) {
	IDOS->Printf(PROGNAME" v%ld.%ld - Copyright (c) 2011 Fredrik Wikstrom\n", VERSION, REVISION);
}

static void show_usage (void) {
	IDOS->Printf("Usage: "PROGNAME" <command> <.vdi file> [options]\n");
	IDOS->Printf("Commands:\n");
	IDOS->Printf("  create          Create a new .vdi file\n");
	IDOS->Printf("  info            List information about .vdi file\n");
	IDOS->Printf("  compact         Compact a dynamic .vdi file\n");
	IDOS->Printf("  convertfromraw  Convert raw (.hdf) to .vdi file\n");
	IDOS->Printf("  converttoraw    Convert .vdi to raw (.hdf) file\n");
	IDOS->Printf("Command syntax:\n");
	IDOS->Printf("  "PROGNAME" create <.vdi file> [fixed|dynamic] <size>\n");
	IDOS->Printf("  "PROGNAME" info <.vdi file>\n");
	IDOS->Printf("  "PROGNAME" compact <.vdi file>\n");
	IDOS->Printf("  "PROGNAME" convertfromraw <.vdi file> <raw file> [fixed|dynamic]\n");
	IDOS->Printf("  "PROGNAME" converttoraw <.vdi file> <raw file>\n");
}

static int vdi_create (int argc, char **argv) {
	BOOL dynamic = FALSE;
	int32 gigabytes = 0;
	VDIFILE *vdi = NULL;
	uint64 disk_size;
	uint64 file_size;
	vdi_header_t *header;
	uint32 block;
	int retval = RETURN_FAIL;

	if (argc == 4) {
		if (IDOS->StrToLong(argv[3], &gigabytes) == -1) {
			show_usage();
			goto error;
		}
	} else if (argc == 5) {
		if (!IUtility->Stricmp(argv[3], "fixed")) {
			dynamic = FALSE;
		} else if (!IUtility->Stricmp(argv[3], "dynamic")) {
			dynamic = TRUE;
		} else {
			show_usage();
			goto error;
		}
		if (IDOS->StrToLong(argv[4], &gigabytes) == -1) {
			show_usage();
			goto error;
		}
	} else {
		show_usage();
		goto error;
	}
	if (gigabytes <= 0) {
		IDOS->FPrintf(StdErr, "VDI size must be > 0 GiB\n");
		goto error;
	}

	vdi = open_vdi(argv[2], MODE_NEWFILE);
	if (!vdi) {
		goto error;
	}

	vdi->block_size = 1UL << 20;
	vdi->total_blocks = gigabytes << 10;
	vdi->offset_blocks = sizeof(vdi_header_t);
	vdi->offset_data = vdi->offset_blocks + ((vdi->total_blocks * sizeof(uint32) + 511UL) & ~511UL);

	file_size = vdi->offset_data;
	disk_size = (uint64)vdi->total_blocks * (uint64)vdi->block_size;
	if (!dynamic) file_size += disk_size;
	if (!IDOS->ChangeFileSize(vdi->file, file_size, OFFSET_BEGINNING)) {
		IDOS->PrintFault(IDOS->IoErr(), vdi->filename);
		IDOS->FPrintf(StdErr, "%s: failed to allocate disk space for vdi file\n", vdi->filename);
		goto error;
	}

	header = vdi->header;
	IUtility->Strlcpy(header->signature, "<<< Sun VirtualBox Disk Image >>>\n", sizeof(header->signature));
	wle32(&header->sig32, VDI_MAGIC);
	wle32(&header->ver32, VDI_VERSION_1_1);
	wle32(&header->hdr_size, 0x190);
	wle32(&header->img_type, dynamic ? DYNAMIC_VDI : FIXED_VDI);
	wle32(&header->offset_blocks, vdi->offset_blocks);
	wle32(&header->offset_data, vdi->offset_data);
	wle32(&header->sector_size, 512);
	wle64(&header->disk_size, disk_size);
	wle32(&header->block_size, vdi->block_size);
	wle32(&header->blocks_in_hdd, vdi->total_blocks);
	wle32(&header->blocks_allocated, dynamic ? 0 : vdi->total_blocks);
	if (!write_header(vdi)) {
		goto error;
	}

	vdi->block_map = IExec->AllocVecTagList(vdi->total_blocks * sizeof(uint32), NULL);
	if (!vdi->block_map) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}
	if (dynamic) {
		for (block = 0; block < vdi->total_blocks; block++) {
			vdi->block_map[block] = -1UL;
		}
	} else {
		for (block = 0; block < vdi->total_blocks; block++) {
			wle32(&vdi->block_map[block], block);
		}
	}
	if (!write_block_map(vdi)) {
		goto error;
	}

	retval = RETURN_OK;

error:
	close_vdi(vdi);

	return retval;
}

static int vdi_info (int argc, char **argv) {
	VDIFILE *vdi = NULL;
	vdi_header_t *header;
	uint32 version;
	CONST_STRPTR image_type;
	uint64 disk_size;
	int retval = RETURN_FAIL;

	if (argc != 3) {
		show_usage();
		goto error;
	}

	vdi = open_vdi(argv[2], MODE_OLDFILE);
	if (!vdi) {
		goto error;
	}

	if (!read_header(vdi)) {
		goto error;
	}

	header = vdi->header;
	version = rle32(&header->ver32);
	IDOS->Printf("Version: %ld.%ld\n", version >> 16, version & 0xffff);
	IDOS->Printf("Header size: %ld (0x%lx)\n", rle32(&header->hdr_size), rle32(&header->hdr_size));
	switch (rle32(&header->img_type)) {
		case FIXED_VDI:
			image_type = "fixed";
			break;
		case DYNAMIC_VDI:
			image_type = "dynamic";
			break;
		default:
			image_type = "unknown";
			break;
	}
	IDOS->Printf("Image type: %s\n", image_type);
	IDOS->Printf("Image flags: 0x%08lx\n", rle32(&header->img_flags));
	IDOS->Printf("Description: %s\n", header->description[0] ? header->description : "<none>");
	IDOS->Printf("Offset blocks: %ld (0x%lx)\n", rle32(&header->offset_blocks), rle32(&header->offset_blocks));
	IDOS->Printf("Offset data: %ld (0x%lx)\n", rle32(&header->offset_data), rle32(&header->offset_data));
	IDOS->Printf("Cylinders: %ld\n", rle32(&header->cylinders));
	IDOS->Printf("Heads: %ld\n", rle32(&header->heads));
	IDOS->Printf("Sectors: %ld\n", rle32(&header->sectors));
	IDOS->Printf("Sector size: %ld (0x%lx)\n", rle32(&header->sector_size), rle32(&header->sector_size));
	disk_size = rle64(&header->disk_size);
	IDOS->Printf("Disk size: %ld.%02ld GiB ", (uint32)(disk_size >> 30),
		(uint32)(((disk_size & 0x3fffffffULL) * 100ULL) >> 30));
	if ((disk_size >> 32) > 0) {
		IDOS->Printf("(0x%lx%08lx)\n", (uint32)(disk_size >> 32), (uint32)disk_size);
	} else {
		IDOS->Printf("(0x%lx)\n", (uint32)disk_size);
	}
	IDOS->Printf("Block size: %ld (0x%lx)\n", rle32(&header->block_size), rle32(&header->block_size));
	IDOS->Printf("Block extra data: %ld\n", rle32(&header->block_extra_data));
	IDOS->Printf("Blocks in HDD: %ld\n", rle32(&header->blocks_in_hdd));
	IDOS->Printf("Blocks allocated: %ld\n", rle32(&header->blocks_allocated));

	retval = RETURN_OK;

error:
	close_vdi(vdi);

	return retval;
}

static int vdi_compact (int argc, char **argv) {
	VDIFILE *vdi = NULL;
	uint32 block_map_size;
	uint32 *block_map = NULL;
	uint8 *block_buffer = NULL;
	uint32 blocks_allocated;
	uint32 blocks_found;
	uint32 blocks_to_compact;
	uint32 blocks_in_use;
	uint32 block;
	int retval = RETURN_FAIL;

	if (argc != 3) {
		show_usage();
		goto error;
	}

	vdi = open_vdi(argv[2], MODE_OLDFILE);
	if (!vdi) {
		goto error;
	}

	if (!read_header(vdi)) {
		goto error;
	}

	if (rle32(&vdi->header->ver32) != VDI_VERSION_1_1) {
		IDOS->FPrintf(StdErr, "%s: unsupported vdi version\n", vdi->filename);
		goto error;
	}

	if (rle32(&vdi->header->img_type) != DYNAMIC_VDI) {
		IDOS->FPrintf(StdErr, "%s: not a dynamic vdi file\n", vdi->filename);
		goto error;
	}

	if (!read_block_map(vdi)) {
		goto error;
	}

	block_map_size = vdi->total_blocks * sizeof(uint32);
	block_map = IExec->AllocVecTagList(block_map_size, NULL);
	if (!block_map) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}

	block_buffer = IExec->AllocVecTagList(vdi->block_size, NULL);
	if (!block_buffer) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}

	blocks_allocated = rle32(&vdi->header->blocks_allocated);
	blocks_found = blocks_to_compact = blocks_in_use = 0;
	for (block = 0; block < vdi->total_blocks; block++) {
		display_progress("checking", blocks_found, blocks_allocated);
		if (IDOS->CheckSignal(SIGBREAKF_CTRL_C)) {
			stop_progress();
			IDOS->PrintFault(ERROR_BREAK, PROGNAME);
			goto error;
		}
		block_map[block] = vdi->block_map[block];
		if (block_map[block] != -1UL) {
			blocks_found++;
			if (!read_block(vdi, block_buffer, block)) {
				goto error;
			}
			if (is_zero_block(block_buffer, vdi->block_size)) {
				blocks_to_compact++;
				block_map[block] = -1UL;
			} else {
				blocks_in_use++;
			}
		}
	}
	display_progress("checking", blocks_found, blocks_allocated);
	stop_progress();

	blocks_allocated = blocks_found;
	IDOS->Printf("Blocks allocated: %lu\n", blocks_allocated);
	IDOS->Printf("Blocks to compact: %lu\n", blocks_to_compact);
	IDOS->Printf("Blocks containing data: %lu\n", blocks_in_use);

	if (blocks_to_compact > 0) {
		uint32 blocks_to_move;
		uint32 blocks_moved;
		uint32 src_block;
		uint32 dst_block;
		uint64 file_size;

		blocks_to_move = blocks_moved = 0;
		for (block = 0; block < blocks_in_use; block++) {
			if (block_map[block] != vdi->block_map[block]) {
				blocks_to_move++;
			}
		}

		src_block = blocks_in_use;
		for (dst_block = 0; dst_block < blocks_in_use; dst_block++) {
			display_progress("compacting", blocks_moved, blocks_to_move);
			if (IDOS->CheckSignal(SIGBREAKF_CTRL_C)) {
				stop_progress();
				IDOS->PrintFault(ERROR_BREAK, PROGNAME);
				goto error;
			}
			if (block_map[dst_block] != vdi->block_map[dst_block]) {
				while (vdi->block_map[src_block] == -1UL) src_block++;
				if (!read_block(vdi, block_buffer, src_block)) {
					goto error;
				}
				if (!write_block(vdi, block_buffer, dst_block)) {
					goto error;
				}
				block_map[src_block] = -1UL;
				block_map[dst_block] = vdi->block_map[dst_block];
				blocks_moved++;
			}
		}
		display_progress("compacting", blocks_moved, blocks_to_move);
		stop_progress();

		wle32(&vdi->header->blocks_allocated, blocks_in_use);
		if (!write_header(vdi)) {
			goto error;
		}
		IExec->CopyMem(block_map, vdi->block_map, block_map_size);
		if (!write_block_map(vdi)) {
			goto error;
		}
		file_size = (uint64)vdi->offset_data + (uint64)blocks_in_use * (uint64)vdi->block_size;
		if (!IDOS->ChangeFileSize(vdi->file, file_size, OFFSET_BEGINNING)) {
			IDOS->PrintFault(IDOS->IoErr(), vdi->filename);
			IDOS->FPrintf(StdErr, "%s: failed to resize vdi file\n", vdi->filename);
			goto error;
		}
		IDOS->Printf("compacting finished\n");
	}

	retval = RETURN_OK;

error:
	IExec->FreeVec(block_buffer);
	IExec->FreeVec(block_map);
	close_vdi(vdi);

	return retval;
}

static int vdi_convertfromraw (int argc, char **argv) {
	BOOL dynamic = FALSE;
	VDIFILE *vdi = NULL;
	RAWFILE *raw = NULL;
	uint64 disk_size;
	uint32 blocks_allocated;
	uint64 file_size;
	vdi_header_t *header;
	uint8 *block_buffer = NULL;
	uint32 block;
	uint32 read_size;
	int32 bytes;
	int retval = RETURN_FAIL;

	if (argc == 5) {
		if (!IUtility->Stricmp(argv[4], "fixed")) {
			dynamic = FALSE;
		} else if (!IUtility->Stricmp(argv[4], "dynamic")) {
			dynamic = TRUE;
		} else {
			show_usage();
			goto error;
		}
	} else if (argc != 4) {
		show_usage();
		goto error;
	}

	raw = open_raw(argv[3], MODE_OLDFILE);
	if (!raw) {
		goto error;
	}

	disk_size = IDOS->GetFileSize(raw->file);
	if (disk_size == -1ULL) {
		IDOS->PrintFault(IDOS->IoErr(), raw->filename);
		goto error;
	}

	vdi = open_vdi(argv[2], MODE_NEWFILE);
	if (!vdi) {
		goto error;
	}

	vdi->block_size = 1UL << 20;
	vdi->total_blocks = (disk_size + (uint64)vdi->block_size - 1ULL) / (uint64)vdi->block_size;
	vdi->offset_blocks = sizeof(vdi_header_t);
	vdi->offset_data = vdi->offset_blocks + ((vdi->total_blocks * sizeof(uint32) + 511UL) & ~511UL);

	blocks_allocated = dynamic ? 0 : vdi->total_blocks;
	file_size = vdi->offset_data;
	if (!dynamic) file_size += (uint64)vdi->total_blocks * (uint64)vdi->block_size;
	if (!IDOS->ChangeFileSize(vdi->file, file_size, OFFSET_BEGINNING)) {
		IDOS->PrintFault(IDOS->IoErr(), vdi->filename);
		IDOS->FPrintf(StdErr, "%s: failed to allocate disk space for vdi file\n", vdi->filename);
		goto error;
	}

	header = vdi->header;
	IUtility->Strlcpy(header->signature, "<<< Sun VirtualBox Disk Image >>>\n", sizeof(header->signature));
	wle32(&header->sig32, VDI_MAGIC);
	wle32(&header->ver32, VDI_VERSION_1_1);
	wle32(&header->hdr_size, 0x190);
	wle32(&header->img_type, dynamic ? DYNAMIC_VDI : FIXED_VDI);
	wle32(&header->offset_blocks, vdi->offset_blocks);
	wle32(&header->offset_data, vdi->offset_data);
	wle32(&header->sector_size, 512);
	wle64(&header->disk_size, disk_size);
	wle32(&header->block_size, vdi->block_size);
	wle32(&header->blocks_in_hdd, vdi->total_blocks);
	wle32(&header->blocks_allocated, blocks_allocated);
	if (!write_header(vdi)) {
		goto error;
	}

	vdi->block_map = IExec->AllocVecTagList(vdi->total_blocks * sizeof(uint32), NULL);
	if (!vdi->block_map) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}
	if (dynamic) {
		for (block = 0; block < vdi->total_blocks; block++) {
			vdi->block_map[block] = -1UL;
		}
	} else {
		for (block = 0; block < vdi->total_blocks; block++) {
			wle32(&vdi->block_map[block], block);
		}
	}
	if (!write_block_map(vdi)) {
		goto error;
	}

	block_buffer = IExec->AllocVecTagList(vdi->block_size, NULL);
	if (!block_buffer) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}

	read_size = vdi->block_size;
	for (block = 0; block < vdi->total_blocks; block++) {
		display_progress("converting", block, vdi->total_blocks);
		if (IDOS->CheckSignal(SIGBREAKF_CTRL_C)) {
			stop_progress();
			IDOS->PrintFault(ERROR_BREAK, PROGNAME);
			goto error;
		}
		if (disk_size < vdi->block_size) {
			read_size = disk_size;
		}
		if ((bytes = IDOS->Read(raw->file, block_buffer, read_size))
			!= read_size)
		{
			stop_progress();
			if (bytes == -1) {
				IDOS->PrintFault(IDOS->IoErr(), raw->filename);
			}
			IDOS->FPrintf(StdErr, "%s: failed to read block %lu\n", raw->filename, block);
			goto error;
		}
		if (read_size < vdi->block_size) {
			IUtility->ClearMem(block_buffer + read_size, vdi->block_size - read_size);
		}
		if (dynamic) {
			if (!is_zero_block(block_buffer, vdi->block_size)) {
				wle32(&vdi->block_map[block], blocks_allocated);
				if (!write_block(vdi, block_buffer, block)) {
					goto error;
				}
				blocks_allocated++;
			}
		} else {
			if (!write_block(vdi, block_buffer, block)) {
				goto error;
			}
		}
		disk_size -= read_size;
	}
	display_progress("converting", block, vdi->total_blocks);
	stop_progress();

	if (dynamic) {
		wle32(&header->blocks_allocated, blocks_allocated);
		if (!write_header(vdi)) {
			goto error;
		}
		if (!write_block_map(vdi)) {
			goto error;
		}
	}

	IDOS->Printf("conversion finished\n");

	retval = RETURN_OK;

error:
	IExec->FreeVec(block_buffer);
	close_vdi(vdi);
	close_raw(raw);

	return retval;
}

static int vdi_converttoraw (int argc, char **argv) {
	VDIFILE *vdi = NULL;
	RAWFILE *raw = NULL;
	uint64 disk_size;
	uint8 *block_buffer = NULL;
	uint32 block;
	uint32 write_size;
	int32 bytes;
	int retval = RETURN_FAIL;

	if (argc != 4) {
		show_usage();
		goto error;
	}

	vdi = open_vdi(argv[2], MODE_OLDFILE);
	if (!vdi) {
		goto error;
	}

	if (!read_header(vdi)) {
		goto error;
	}
	if (!read_block_map(vdi)) {
		goto error;
	}
	disk_size = rle64(&vdi->header->disk_size);

	raw = open_raw(argv[3], MODE_NEWFILE);
	if (!raw) {
		goto error;
	}

	if (!IDOS->ChangeFileSize(raw->file, disk_size, OFFSET_BEGINNING)) {
		IDOS->PrintFault(IDOS->IoErr(), raw->filename);
		IDOS->FPrintf(StdErr, "%s: failed to allocate disk space for raw file\n", raw->filename);
		goto error;
	}

	block_buffer = IExec->AllocVecTagList(vdi->block_size, NULL);
	if (!block_buffer) {
		IDOS->PrintFault(ERROR_NO_FREE_STORE, PROGNAME);
		goto error;
	}

	write_size = vdi->block_size;
	for (block = 0; block < vdi->total_blocks; block++) {
		display_progress("converting", block, vdi->total_blocks);
		if (IDOS->CheckSignal(SIGBREAKF_CTRL_C)) {
			stop_progress();
			IDOS->PrintFault(ERROR_BREAK, PROGNAME);
			goto error;
		}
		if (vdi->block_map[block] == -1UL) {
			IUtility->ClearMem(block_buffer, vdi->block_size);
		} else {
			if (!read_block(vdi, block_buffer, block)) {
				goto error;
			}
		}
		if (disk_size < vdi->block_size) {
			write_size = disk_size;
		}
		if ((bytes = IDOS->Write(raw->file, block_buffer, write_size))
			!= write_size)
		{
			stop_progress();
			IDOS->PrintFault(IDOS->IoErr(), raw->filename);
			IDOS->FPrintf(StdErr, "%s: failed to write block %lu\n", raw->filename, block);
			goto error;
		}
		disk_size -= write_size;
	}
	display_progress("converting", block, vdi->total_blocks);
	stop_progress();

	IDOS->Printf("conversion finished\n");

	retval = RETURN_OK;

error:
	IExec->FreeVec(block_buffer);
	close_vdi(vdi);
	close_raw(raw);

	return retval;
}

static BOOL is_zero_block (const uint8 *buffer, uint32 block_size) {
	const uint32 *ptr = (const uint32 *)buffer;
	while (block_size > 0) {
		if ((ptr[0] | ptr[1]) | (ptr[2] | ptr[3])) {
			return FALSE;
		}
		ptr += 4;
		block_size -= 16;
	}
	return TRUE;
}
