/*
 * Copyright 2003-2009, Haiku, Inc. All Rights Reserved.
 * Distributed under the terms of the MIT License.
 *
 * Authors:
 *		Bryce Groff, brycegroff@gmail.com
 */

#include "PartitionMapWriter.h"

#include <errno.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>

#include <new>

#ifndef _USER_MODE
#include <debug.h>
#endif

#ifndef _USER_MODE
#	include <KernelExport.h>
#endif

#include "PartitionMap.h"

using std::nothrow;


#define TRACE_ENABLED
#ifdef TRACE_ENABLED
#	ifdef _USER_MODE
#		define TRACE(x) printf x
#	else
#		define TRACE(x) dprintf x
#	endif
#endif


#if defined(__i386__) || defined(__x86_64__)
#	ifndef _USER_MODE
#		define MBR_HEADER "MBR.h"
#		include MBR_HEADER
#	endif
#endif


bool
check_logical_location(const LogicalPartition* child,
	const PrimaryPartition* parent)
{
	if (child->PartitionTableOffset() % child->BlockSize() != 0) {
		TRACE(("check_logical_location() - PartitionTableOffset: %" B_PRId64 " "
			"not a multiple of media's block size: %" B_PRId32 "\n",
			child->PartitionTableOffset(), child->BlockSize()));
		return false;
	}
	if (child->Offset() % child->BlockSize() != 0) {
		TRACE(("check_logical_location() - Parition offset: %" B_PRId64 " "
			"is not a multiple of block size: %" B_PRId32 "\n", child->Offset(),
			child->BlockSize()));
		return false;
	}
	if (child->Size() % child->BlockSize() != 0) {
		TRACE(("check_logical_location() - Size: (%" B_PRId64 ") is not a "
			"multiple of block size: (%" B_PRId32 ")\n", child->Size(),
			child->BlockSize()));
		return false;
	}
	if (child->PartitionTableOffset() < parent->Offset()
		|| child->PartitionTableOffset() >= parent->Offset()
		+ parent->Size()) {
		TRACE(("check_logical_location() - Partition table: (%" B_PRId64 ") not"
			" within extended partition (start: %" B_PRId64 "), (end: "
			"%" B_PRId64 ")\n", child->PartitionTableOffset(), parent->Offset(),
			parent->Offset() + parent->Size()));
		return false;
	}
	if (child->Offset() + child->Size() > parent->Offset() + parent->Size()) {
		TRACE(("check_logical_location() - logical paritition does not lie "
			"within extended partition\n"));
		return false;
	}
	return true;
}


PartitionMapWriter::PartitionMapWriter(int deviceFD, uint32 blockSize)
	:
	fDeviceFD(deviceFD),
	fBlockSize(blockSize)
{
}


PartitionMapWriter::~PartitionMapWriter()
{
}


status_t
PartitionMapWriter::WriteMBR(const PartitionMap* map, bool writeBootCode)
{
	if (map == NULL)
		return B_BAD_VALUE;

	partition_table partitionTable;
	status_t error = _ReadBlock(0, partitionTable);
	if (error != B_OK)
		return error;
#ifdef MBR_HEADER
	if (writeBootCode) {
		// the boot code must be small enough to fit in the code area
		STATIC_ASSERT(kMBRSize <= sizeof(partitionTable.code_area));
		partitionTable.clear_code_area();
		partitionTable.fill_code_area(kMBR, kMBRSize);
	}
#endif

	partitionTable.signature = kPartitionTableSectorSignature;

	for (int i = 0; i < 4; i++) {
		partition_descriptor* descriptor = &partitionTable.table[i];
		const PrimaryPartition* partition = map->PrimaryPartitionAt(i);

		partition->GetPartitionDescriptor(descriptor);
	}

	error = _WriteBlock(0, partitionTable);
	return error;
}


status_t
PartitionMapWriter::WriteLogical(const LogicalPartition* logical,
	const PrimaryPartition* primary, bool clearCode)
{
	if (logical == NULL || primary == NULL)
		return B_BAD_VALUE;

	if (!check_logical_location(logical, primary))
		return B_BAD_DATA;

	partition_table partitionTable;
	if (clearCode) {
		partitionTable.clear_code_area();
	} else {
		status_t error = _ReadBlock(logical->PartitionTableOffset(),
			partitionTable);
		if (error != B_OK)
			return error;
	}

	partitionTable.signature = kPartitionTableSectorSignature;

	partition_descriptor* descriptor = &partitionTable.table[0];
	logical->GetPartitionDescriptor(descriptor);

	descriptor = &partitionTable.table[1];
	if (logical->Next() != NULL)
		logical->Next()->GetPartitionDescriptor(descriptor, true);
	else
		memset(descriptor, 0, sizeof(partition_descriptor));

	// last two descriptors are empty
	for (int32 i = 2; i < 4; i++) {
		descriptor = &partitionTable.table[i];
		memset(descriptor, 0, sizeof(partition_descriptor));
	}

	status_t error = _WriteBlock(logical->PartitionTableOffset(),
		partitionTable);
	return error;
}


status_t
PartitionMapWriter::WriteExtendedHead(const LogicalPartition* logical,
	const PrimaryPartition* primary, bool clearCode)
{
	if (primary == NULL)
		return B_BAD_VALUE;

	partition_table partitionTable;
	if (clearCode) {
		partitionTable.clear_code_area();
	} else {
		status_t error = _ReadBlock(primary->Offset(), partitionTable);
		if (error != B_OK)
			return error;
	}

	partitionTable.signature = kPartitionTableSectorSignature;
	partition_descriptor* descriptor;
	if (logical == NULL) {
		for (int32 i = 0; i < 4; i++) {
			descriptor = &partitionTable.table[i];
			memset(descriptor, 0, sizeof(partition_descriptor));
		}
	} else {
		LogicalPartition partition;
		partition.SetPartitionTableOffset(primary->Offset());
		partition.SetBlockSize(logical->BlockSize());
		partition.SetOffset(logical->Offset());
		partition.SetSize(logical->Size());
		partition.SetType(logical->Type());

		// set the logicals partition table to the correct location
		descriptor = &partitionTable.table[0];
		partition.GetPartitionDescriptor(descriptor);

		descriptor = &partitionTable.table[1];
		LogicalPartition* next = logical->Next();
		if (next != NULL)
			next->GetPartitionDescriptor(descriptor, true);
		else
			memset(descriptor, 0, sizeof(partition_descriptor));

		// last two descriptors are empty
		for (int32 i = 2; i < 4; i++) {
			descriptor = &partitionTable.table[i];
			memset(descriptor, 0, sizeof(partition_descriptor));
		}
	}

	status_t error = _WriteBlock(primary->Offset(), partitionTable);
	if (error != B_OK)
		return error;

	return B_OK;
}



status_t
PartitionMapWriter::ClearExtendedHead(const PrimaryPartition* primary)
{
	if (primary == NULL)
		return B_BAD_VALUE;

	partition_table partitionTable;
	partitionTable.clear_code_area();
	partitionTable.signature = kPartitionTableSectorSignature;

	partition_descriptor* descriptor;
	for (int32 i = 0; i < 4; i++) {
		descriptor = &partitionTable.table[i];
		memset(descriptor, 0, sizeof(partition_descriptor));
	}

	status_t error = _WriteBlock(primary->Offset(), partitionTable);
	if (error != B_OK)
		return error;

	return B_OK;
}


status_t
PartitionMapWriter::_ReadBlock(off_t partitionOffset,
	partition_table& partitionTable)
{
	if (partitionOffset < 0)
		return B_BAD_VALUE;
	// TODO: If fBlockSize > sizeof(partition_table) then stop/read NULL after
	if (read_pos(fDeviceFD, partitionOffset, &partitionTable,
		sizeof(partitionTable)) != sizeof(partitionTable)) {
		status_t error = errno;
		if (error == B_OK)
			error = B_IO_ERROR;

		return error;
	}

	return B_OK;
}


status_t
PartitionMapWriter::_WriteBlock(off_t partitionOffset,
	const partition_table& partitionTable)
{
	if (partitionOffset < 0)
		return B_BAD_VALUE;
	// TODO: maybe clear the rest of the block if
	// fBlockSize > sizeof(partition_table)?
	if (write_pos(fDeviceFD, partitionOffset, &partitionTable,
		sizeof(partitionTable)) != sizeof(partitionTable)) {
		status_t error = errno;
		if (error == B_OK)
			error = B_IO_ERROR;

		return error;
	}

	return B_OK;
}