/*
 * Copyright 2010 Michael Lotz, mmlr@mlotz.ch
 * Copyright 2011-2017 Haiku, Inc. All rights reserved.
 * All rights reserved. Distributed under the terms of the MIT License.
 *
 * Authors:
 * 		Michael Lotz <mmlr@mlotz.ch>
 * 		Alexander von Gluck IV <kallisti5@unixzen.com>
 */


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

#define EXIT_FAILURE 1


//#define DEBUG_ANYBOOT
#ifdef DEBUG_ANYBOOT
#   define TRACE(x...) printf("anyboot: " x)
#else
#   define TRACE(x...) ;
#endif


static uint8_t *sCopyBuffer = NULL;
static size_t sCopyBufferSize = 2 * 1024 * 1024;
static const size_t kBlockSize = 512;


// Haiku Anyboot Image:
//   (MBR Table + Boot Sector)
//   ISO (Small Haiku ISO9660)
//   First Partition (Haiku OS Image, BFS)
//   Second Partition (EFI Loader, FAT)
//   Third Partition (EFI Mac, HFS) <not implemented>


static void
print_usage(bool error)
{
	printf("\n");
	printf("usage: anyboot [-b BIOS-Loader] [-e EFI-Filesystem] <isoFile> <imageFile> <outputFile>\n");
	printf("       -b, --bios-loader <file>     Legacy BIOS bootloader\n");
	printf("       -e, --efi-filesystem <file>  EFI filesystem\n");
	exit(error ? EXIT_FAILURE : 0);
}


static void
checkError(bool failed, const char *message)
{
	if (!failed)
		return;

	printf("%s: %s\n", message, strerror(errno));
	free(sCopyBuffer);
	exit(1);
}


#if 0
static void
chsAddressFor(uint32_t offset, uint8_t *address, uint32_t sectorsPerTrack,
	uint32_t headsPerCylinder)
{
	if (offset >= 1024 * sectorsPerTrack * headsPerCylinder) {
		// not encodable, force LBA
		address[0] = 0xff;
		address[1] = 0xff;
		address[2] = 0xff;
		return;
	}

	uint32_t cylinders = 0;
	uint32_t heads = 0;
	uint32_t sectors = 0;

	uint32_t temp = 0;
	while (temp * sectorsPerTrack * headsPerCylinder <= offset)
		cylinders = temp++;

	offset -= (sectorsPerTrack * headsPerCylinder * cylinders);

	temp = 0;
	while (temp * sectorsPerTrack <= offset)
		heads = temp++;

	sectors = offset - (sectorsPerTrack * heads) + 1;

	address[0] = heads;
	address[1] = sectors;
	address[1] |= (cylinders >> 2) & 0xc0;
	address[2] = cylinders;
}
#endif


static void
createPartition(int handle, int index, bool active, uint8_t type,
	uint32_t offset, uint32_t size)
{
	uint8_t bootable = active ? 0x80 : 0x0;
	uint8_t partition[16] = {
		bootable,				// bootable
		0xff, 0xff, 0xff,		// CHS first block (default to LBA)
		type,					// partition type
		0xff, 0xff, 0xff,		// CHS last block (default to LBA)
		0x00, 0x00, 0x00, 0x00,	// imageOffset in blocks (written below)
		0x00, 0x00, 0x00, 0x00	// imageSize in blocks (written below)
	};

	// fill in LBA values
	uint32_t partitionOffset = (uint32_t)(offset / kBlockSize);
	((uint32_t *)partition)[2] = partitionOffset;
	((uint32_t *)partition)[3] = (uint32_t)(size / kBlockSize);

	TRACE("%s: #%d %c bytes: %u-%u, sectors: %u-%u\n", __func__, index,
		active ? 'b' : '-', offset, offset + size, partitionOffset,
		partitionOffset + uint32_t(size / kBlockSize));
#if 0
	// while this should basically work, it makes the boot code needlessly
	// use chs which has a high potential of failure due to the geometry
	// being unknown beforehand (a fixed geometry would be used here).

	uint32_t sectorsPerTrack = 63;
	uint32_t headsPerCylinder = 255;

	// fill in CHS values
	chsAddressFor(partitionOffset, &partition[1], sectorsPerTrack,
		headsPerCylinder);
	chsAddressFor(partitionOffset + (uint32_t)(size / kBlockSize),
		&partition[5], sectorsPerTrack, headsPerCylinder);
#endif

	ssize_t written = pwrite(handle, partition, 16, 512 - 2 - 16 * (4 - index));
	checkError(written != 16, "failed to write partition entry");

	if (active) {
		// make it bootable
		written = pwrite(handle, &partitionOffset, 4, offset + 512 - 2 - 4);
		checkError(written != 4, "failed to make image bootable");
	}
	return;
}


static int
copyLoop(int from, int to, off_t position)
{
	while (true) {
		ssize_t copyLength = read(from, sCopyBuffer, sCopyBufferSize);
		if (copyLength <= 0)
			return copyLength;

		ssize_t written = pwrite(to, sCopyBuffer, copyLength, position);
		if (written != copyLength) {
			if (written < 0)
				return written;
			else
				return -1;
		}

		position += copyLength;
	}
}


int
main(int argc, char *argv[])
{
	const char *biosFile = NULL;
	const char *efiFile = NULL;
	const char *isoFile = NULL;
	const char *imageFile = NULL;
	const char *outputFile = NULL;

	while (1) {
		int c;
		static struct option long_options[] = {
			{"bios-loader", required_argument, 0, 'b'},
			{"efi-loader", required_argument, 0, 'e'},
			{"help", no_argument, 0, 'h'},
			{0, 0, 0, 0}
		};

		opterr = 1; /* don't print errors */
		c = getopt_long(argc, argv, "+hb:e:", long_options, NULL);

		if (c == -1)
			break;

		switch (c) {
			case 'h':
				print_usage(false);
				break;
			case 'b':
				biosFile = optarg;
				break;
			case 'e':
				efiFile = optarg;
				break;
			default:
				print_usage(true);
		}
	}

	if ((argc - optind) != 3)
		print_usage(true);

	for (int index = optind; index < argc; index++) {
		if (isoFile == NULL)
			isoFile = argv[index];
		else if (imageFile == NULL)
			imageFile = argv[index];
		else if (outputFile == NULL)
			outputFile = argv[index];
	}

	sCopyBuffer = (uint8_t *)malloc(sCopyBufferSize);
	checkError(sCopyBuffer == NULL, "no memory for copy buffer");

	int outputFileHandle = open(outputFile, O_WRONLY | O_TRUNC | O_CREAT,
		S_IRUSR | S_IWUSR);
	checkError(outputFileHandle < 0, "failed to open output file");

	int isoFileHandle = open(isoFile, O_RDONLY);
	checkError(isoFileHandle < 0, "failed to open ISO file");

	struct stat stat;
	int result = fstat(isoFileHandle, &stat);
	checkError(result != 0, "failed to stat ISO file");
	off_t isoSize = stat.st_size;

	int biosFileHandle = -1;
	if (biosFile != NULL) {
		biosFileHandle = open(biosFile, O_RDONLY);
		checkError(biosFileHandle < 0, "failed to open BIOS bootloader file");
	}

	int efiFileHandle = -1;
	off_t efiSize = 0;
	if (efiFile != NULL) {
		efiFileHandle = open(efiFile, O_RDONLY);
		checkError(efiFileHandle < 0, "failed to open EFI bootloader file");

		result = fstat(efiFileHandle, &stat);
		checkError(result != 0, "failed to stat EFI filesystem image");
		efiSize = stat.st_size;
	}

	int imageFileHandle = open(imageFile, O_RDONLY);
	checkError(imageFileHandle < 0, "failed to open image file");

	result = fstat(imageFileHandle, &stat);
	checkError(result != 0, "failed to stat image file");
	off_t imageSize = stat.st_size;

	result = copyLoop(isoFileHandle, outputFileHandle, 0);
	checkError(result != 0, "failed to copy iso file to output");

	// isoSize rounded to the next full megabyte
	off_t alignment = 1 * 1024 * 1024 - 1;
	off_t imageOffset = (isoSize + alignment) & ~alignment;

	result = copyLoop(imageFileHandle, outputFileHandle, imageOffset);
	checkError(result != 0, "failed to copy image file to output");

	if (biosFileHandle >= 0) {
		result = copyLoop(biosFileHandle, outputFileHandle, 0);
		checkError(result != 0, "failed to copy BIOS bootloader to output");
	}

	// Haiku Image Partition
	alignment = kBlockSize - 1;
	imageSize = (imageSize + alignment) & ~alignment;
	createPartition(outputFileHandle, 0, true, 0xeb, imageOffset, imageSize);

	// Optional EFI Filesystem
	if (efiFile != NULL) {
		off_t efiOffset = (imageOffset + imageSize + alignment) & ~alignment;
		efiSize = (efiSize + alignment) & ~alignment;
		result = copyLoop(efiFileHandle, outputFileHandle, efiOffset);
		checkError(result != 0, "failed to copy EFI filesystem image to output");
		createPartition(outputFileHandle, 1, false, 0xef, efiOffset, efiSize);
	}

	// TODO: MacEFI (HFS) maybe someday

	free(sCopyBuffer);
	return 0;
}