/*
 * Copyright 2012 Aleksas Pantechovskis, <alexp.frl@gmail.com>
 * All rights reserved. Distributed under the terms of the MIT License.
 */

#include <fstream>
#include <iostream>
#include <sstream>
#include <string.h>
#include <string>

#include <Debug.h>
#include <DiskDevice.h>
#include <DiskDeviceRoster.h>
#include <Path.h>
#include <Volume.h>
#include <VolumeRoster.h>

#include "MBR.h"


using namespace std;

const char* kUsageMessage = \
	"Usage: writembr [ device ] \n"
	"#\tRewrites the MBR for the specified device.\n"
	"#\tIf no device is specified, the boot device is used.\n"
	"#\t--help shows this usage message\n";


int
main(int argc, char** argv)
{
	if ((argc == 2 && strcmp(argv[1], "--help") == 0) || argc > 2) {
		cerr << kUsageMessage;
		return B_ERROR;
	}

	BPath device;

	if (argc == 2)
		// user specified device for rewriting
		device.SetTo(argv[1]);

	else if (argc == 1) {
		// no parameters specified, rewrite boot device
		BVolumeRoster volumeRoster;
		BVolume bootVolume;
		if (volumeRoster.GetBootVolume(&bootVolume) != B_OK) {
			cerr << "Can not find boot device" << endl;
			return B_ERROR;
		}

		BDiskDeviceRoster roster;
		BDiskDevice bootDevice;
		if(roster.FindPartitionByVolume(bootVolume, &bootDevice, NULL) != B_OK) {
			cerr << "Can not find boot device" << endl;
			return B_ERROR;
		}

		bootDevice.GetPath(&device);
	}


	if (strcmp(device.Leaf(), "raw") != 0) {
		cerr << device.Path() << " is not a raw device" << endl;
		return B_ERROR;
	}

	fstream fs;
	fs.open(device.Path(), fstream::in | fstream::out | fstream::binary);
	if (!fs.is_open()) {
		cerr << "Can't open " << device.Path() << endl;
		return B_ERROR;
	}

	STATIC_ASSERT(kMBRSize == 512);

	unsigned char MBR[kMBRSize];
	fs.read((char*)MBR, kMBRSize);
	if (fs.fail() || fs.gcount() < (off_t)kMBRSize ) {
		cerr << "Cannot read " << kMBRSize
			<< " bytes from " << device.Path() << endl;
		fs.close();
		return B_ERROR;
	}

	// update only the code area and the MBR signature
	memcpy(MBR, kMBR, 0x1be);
	MBR[0x1FE] = kMBR[0x1FE];
	MBR[0x1FF] = kMBR[0x1FF];

	cerr << "About to overwrite the MBR boot code on " << device.Path()
		<< "\nThis may disable any partition managers you have installed.\n"
		<< "Are you sure you want to continue?\nyes/[no]: ";

	string choice;
	getline(cin, choice, '\n');
	if (choice == "no" || choice == "" || choice != "yes") {
		cerr << "MBR was NOT written" << endl;
		fs.close();
		return B_ERROR;
	}
			
	cerr << "Rewriting MBR for " << device.Path() << endl;
		
	fs.seekg(0, ios::beg);
	fs.write((char*)MBR, kMBRSize);
	if (fs.fail()) {
		cerr << "Cannot write " << kMBRSize
			<< " bytes to " << device.Path() << endl;
		fs.close();
		return B_ERROR;
	}

	fs.close();

	cerr << "MBR was written OK" << endl;
	return B_OK;
}