/*
 * Copyright 2017, Chế Vũ Gia Hy, cvghy116@gmail.com.
 * This file may be used under the terms of the MIT License.
 */


#include "ExtentAllocator.h"
#include "Chunk.h"


CachedExtent*
CachedExtent::Create(uint64 offset, uint64 length, uint64 flags)
{
	CachedExtent* self = new(std::nothrow) CachedExtent();
	if (self == NULL)
		return NULL;

	self->offset = offset;
	self->length = length;
	self->refCount = 0;
	if ((flags & BTRFS_EXTENT_FLAG_ALLOCATED) != 0)
		self->refCount = 1;
	self->flags = flags;
	self->item = NULL;
	return self;
}


void
CachedExtent::Delete()
{
	if (item != NULL)
		free(item);
	item = NULL;
	delete this;
}


bool
CachedExtent::IsAllocated() const
{
	return (flags & BTRFS_EXTENT_FLAG_ALLOCATED) != 0;
}


bool
CachedExtent::IsData() const
{
	return (flags & BTRFS_EXTENT_FLAG_DATA) != 0;
}


void
CachedExtent::Info() const
{
	const char* extentType = "allocated tree block";
	if (IsAllocated() == false && IsData() == false)
		extentType = "free tree block";
	else if (IsAllocated() == false && IsData() == true)
		extentType = "free data extent";
	else if (IsAllocated() == true && IsData() == true)
		extentType = "allocated data extent";

	TRACE("%s at %" B_PRIu64 " length %" B_PRIu64 " refCount %i\n",
		extentType, offset, length, refCount);
}


// ExtentTree Implementation


CachedExtentTree::CachedExtentTree(const TreeDefinition& definition)
	:
	AVLTree<TreeDefinition>(definition)
{
}


CachedExtentTree::~CachedExtentTree()
{
	Delete();
}


/* Find extent that cover or after "offset" and has length >= "size"
 * it must also satisfy the condition "type".
 */
status_t
CachedExtentTree::FindNext(CachedExtent** chosen, uint64 offset, uint64 size,
	uint64 type)
{
	CachedExtent* found = Find(offset);
	while (found != NULL && (found->flags != type || found->length < size))
		found = Next(found);

	if (found == NULL)
		return B_ENTRY_NOT_FOUND;
	*chosen = found;
	return B_OK;
}


/* this will insert all free extents that are holes,
 * created by existed allocated extents in the tree
 * from lowerBound to upperBound
 */
status_t
CachedExtentTree::FillFreeExtents(uint64 lowerBound, uint64 upperBound)
{
	CachedExtent* node = FindClosest(lowerBound, false);
	CachedExtent* hole = NULL;
	status_t status = B_OK;
	uint64 flags = node->flags & (~BTRFS_EXTENT_FLAG_ALLOCATED);
	if (lowerBound < node->offset) {
		hole = CachedExtent::Create(lowerBound, node->offset - lowerBound,
			flags);
		status = _AddFreeExtent(hole);
		if (status != B_OK)
			return status;
	}

	CachedExtent* next = NULL;
	while ((next = Next(node)) != NULL && next->End() < upperBound) {
		if (node->End() == next->offset) {
			node = next;
			continue;
		}

		hole = CachedExtent::Create(node->End(), next->offset - node->End(),
			flags);
		status = _AddFreeExtent(hole);
		if (status != B_OK)
			return status;
		node = next;
	}

	// final node should be a right most node
	if (upperBound > node->End()) {
		hole = CachedExtent::Create(node->End(), upperBound - node->End(),
			flags);
		status = _AddFreeExtent(hole);
	}

	return status;
}


void
CachedExtentTree::_RemoveExtent(CachedExtent* node)
{
	node->refCount--;
	if (node->refCount <= 0) {
		Remove(node);
		node->Delete();
	}
}


status_t
CachedExtentTree::_AddAllocatedExtent(CachedExtent* node)
{
	if (node == NULL || node->IsAllocated() == false)
		return B_BAD_VALUE;

	CachedExtent* found = Find(node->offset);
	if (found == NULL) {
		Insert(node);
		return B_OK;
	}

	if ((found->IsData() && !node->IsData())
		|| (!found->IsData() && node->IsData())) {
		// this shouldn't happen as because different type has
		// its different region for locating.
		node->Delete();
		return B_BAD_VALUE;
	}

	if (found->IsAllocated() == false) {
		// split free extents (found)
		if (node->End() > found->End()) {
			// TODO: when to handle this ?
			node->Delete();
			return B_BAD_VALUE;
		}

		// |----- found ------|
		//    |-- node ---|
		uint64 diff = node->offset - found->offset;
		found->offset += diff + node->length;
		found->length -= diff + node->length;
		// diff < 0 couldn't happen because of the Compare function
		if (diff > 0) {
			CachedExtent* leftEmpty = CachedExtent::Create(
				node->offset - diff, diff, found->flags);
			_AddFreeExtent(leftEmpty);
		}

		if (found->length == 0) {
			// free-extent that has no space
			_RemoveExtent(found);
		}

		Insert(node);
		return B_OK;
	}

	if (found->length == node->length) {
		found->refCount++;
	} else {
		// TODO: What should we do in this case ?
		return B_BAD_VALUE;
	}
	node->Delete();

	return B_OK;
}


status_t
CachedExtentTree::_AddFreeExtent(CachedExtent* node)
{
	if (node == NULL || node->IsAllocated() == true)
		return B_BAD_VALUE;

	CachedExtent* found = Find(node->offset);
	if (found == NULL) {
		Insert(node);
		_CombineFreeExtent(node);
		return B_OK;
	}

	if ((found->IsData() && !node->IsData())
		|| (!found->IsData() && node->IsData())) {
		// this shouldn't happen as because different type has
		// its different region for locating.
		node->Delete();
		return B_BAD_VALUE;
	}

	if (found->End() < node->End()) {
		// |---- found-----|
		//      |--- node ------|
		CachedExtent* rightEmpty = CachedExtent::Create(found->End(),
			node->End() - found->End(), node->flags);
		_AddFreeExtent(rightEmpty);
		node->length -= node->End() - found->End();
	}

	if (found->IsAllocated() == true) {
		// free part of the allocated extents(found)
		// |----- found ------|
		//    |-- node ---|
		uint64 diff = node->offset - found->offset;
		found->offset += diff + node->length;
		found->length -= diff + node->length;
		// diff < 0 couldn't happen because of the Compare function
		if (diff > 0){
			CachedExtent* left = CachedExtent::Create(node->offset - diff,
				diff, found->flags);
			_AddAllocatedExtent(left);
		}

		if (found->length == 0)
			_RemoveExtent(found);

		Insert(node);
		_CombineFreeExtent(node);
	}

	return B_OK;
}


status_t
CachedExtentTree::AddExtent(CachedExtent* extent)
{
	if (extent->IsAllocated() == true)
		return _AddAllocatedExtent(extent);
	return _AddFreeExtent(extent);
}


void
CachedExtentTree::_CombineFreeExtent(CachedExtent* node)
{
	// node should be inserted first to call this function,
	// otherwise it will overdelete.
	if (node->IsAllocated() == true)
		return;

	CachedExtent* other = Next(node);
	if (other != NULL) {
		if (node->End() == other->offset && node->flags == other->flags) {
			node->length += other->length;
			_RemoveExtent(other);
		}
	}

	other = Previous(node);
	if (other != NULL) {
		if (other->End() == node->offset && node->flags == other->flags) {
			other->length += node->length;
			_RemoveExtent(node);
		}
	}
}


void
CachedExtentTree::_DumpInOrder(CachedExtent* node) const
{
	if (node == NULL)
		return;
	_DumpInOrder(_GetValue(node->left));
	node->Info();
	_DumpInOrder(_GetValue(node->right));
}


void
CachedExtentTree::DumpInOrder() const
{
	TRACE("\n");
	_DumpInOrder(RootNode());
	TRACE("\n");
}


void
CachedExtentTree::Delete()
{
	_Delete(RootNode());
	Clear();
}


void
CachedExtentTree::_Delete(CachedExtent* node)
{
	if (node == NULL)
		return;
	_Delete(_GetValue(node->left));
	_Delete(_GetValue(node->right));
	node->Delete();
}


// BlockGroup


BlockGroup::BlockGroup(BTree* extentTree)
	:
	fItem(NULL)
{
	fCurrentExtentTree = new(std::nothrow) BTree(extentTree->SystemVolume(),
		extentTree->RootBlock());
}


BlockGroup::~BlockGroup()
{
	if (fItem != NULL) {
		free(fItem);
		fItem = NULL;
	}

	delete fCurrentExtentTree;
	fCurrentExtentTree = NULL;
}


status_t
BlockGroup::SetExtentTree(off_t rootAddress)
{
	status_t status = fCurrentExtentTree->SetRoot(rootAddress, NULL);
	if (status != B_OK)
		return status;

	if (fItem != NULL) {
		// re-allocate BlockGroup;
		uint64 flags = Flags();
		return Initialize(flags);
	}

	return B_OK;
}


/* Initialize BlockGroup that has flag
 */
status_t
BlockGroup::Initialize(uint64 flag)
{
	if (fCurrentExtentTree == NULL)
		return B_NO_MEMORY;

	if (fItem != NULL)
		free(fItem);

	TRACE("BlockGroup::Initialize() find block group has flag: %"
		B_PRIu64 "\n", flag);
	BTree::Path path(fCurrentExtentTree);
	fKey.SetObjectID(0);
	// find first item in extent to get the ObjectID
	// ignore type because block_group is not always the first item
	fKey.SetType(BTRFS_KEY_TYPE_ANY);
	fCurrentExtentTree->FindNext(&path, fKey, NULL);

	fKey.SetType(BTRFS_KEY_TYPE_BLOCKGROUP_ITEM);
	fKey.SetOffset(0);
	status_t status;

	while (true) {
		status = fCurrentExtentTree->FindNext(&path, fKey, (void**)&fItem);
		if (status != B_OK || (Flags() & flag) == flag)
			break;
		fKey.SetObjectID(End());
		fKey.SetOffset(0);
	}

	if (status != B_OK)
		ERROR("BlockGroup::Initialize() cannot find block group\n");

	return status;
}


status_t
BlockGroup::LoadExtent(CachedExtentTree* tree, bool inverse)
{
	if (fItem == NULL)
		return B_NO_INIT;

	uint64 flags = (Flags() & BTRFS_BLOCKGROUP_FLAG_DATA) != 0 ?
		BTRFS_EXTENT_FLAG_DATA : BTRFS_EXTENT_FLAG_TREE_BLOCK;
	if (inverse == false)
		flags |= BTRFS_EXTENT_FLAG_ALLOCATED;

	uint64 start = Start();
	CachedExtent* insert;
	void* data;
	uint64 extentSize = fCurrentExtentTree->SystemVolume()->BlockSize();

	btrfs_key key;
	key.SetType(BTRFS_KEY_TYPE_METADATA_ITEM);
	if ((flags & BTRFS_EXTENT_FLAG_DATA) != 0)
		key.SetType(BTRFS_KEY_TYPE_EXTENT_ITEM);
	key.SetObjectID(start);
	key.SetOffset(0);

	TreeIterator iterator(fCurrentExtentTree, key);
	status_t status;
	while (true) {
		status = iterator.GetNextEntry(&data);
		key = iterator.Key();
		if (status != B_OK) {
			if (key.ObjectID() != Start())
				break;
			// When we couldn't find the item and key has
			// objectid == BlockGroup's objectID, because
			// key's type < BLOCKGROUP_ITEM
			continue;
		}

		if (inverse == false) {
			if ((flags & BTRFS_EXTENT_FLAG_DATA) != 0)
				extentSize = key.Offset();
			insert = CachedExtent::Create(key.ObjectID(), extentSize, flags);
			insert->item = (btrfs_extent*)data;
		} else {
			extentSize = key.ObjectID() - start;
			insert = CachedExtent::Create(start, extentSize, flags);
			free(data);		// free extent doesn't have data;
		}
		_InsertExtent(tree, insert);
		start = key.ObjectID() + extentSize;
	}

	if (inverse == true)
		_InsertExtent(tree, start, End() - start, flags);

	return B_OK;
}


status_t
BlockGroup::_InsertExtent(CachedExtentTree* tree, uint64 start, uint64 length,
	uint64 flags)
{
	CachedExtent* extent = CachedExtent::Create(start, length, flags);
	return _InsertExtent(tree, extent);
}


status_t
BlockGroup::_InsertExtent(CachedExtentTree* tree, CachedExtent* extent)
{
	if (extent->length == 0)
		return B_BAD_VALUE;

	status_t status = tree->AddExtent(extent);
	if (status != B_OK) {
		return status;
	}
	return B_OK;
}


// ExtentAllocator


ExtentAllocator::ExtentAllocator(Volume* volume)
	:
	fVolume(volume),
	fTree(NULL),
	fStart((uint64)-1),
	fEnd(0)
{
	fTree = new(std::nothrow) CachedExtentTree(TreeDefinition());
}


ExtentAllocator::~ExtentAllocator()
{
	delete fTree;
	fTree = NULL;
}


status_t
ExtentAllocator::_LoadExtentTree(uint64 flags)
{
	TRACE("ExtentAllocator::_LoadExtentTree() flags: %" B_PRIu64 "\n", flags);
	BlockGroup blockGroup(fVolume->ExtentTree());
	status_t status = blockGroup.Initialize(flags);
	if (status != B_OK)
		return status;

	for (int i = 0; i < BTRFS_NUM_ROOT_BACKUPS; i++) {
		uint64 extentRootAddr =
			fVolume->SuperBlock().backup_roots[i].ExtentRoot();
		if (extentRootAddr == 0)
			continue;	// new device has 0 root address

		status = blockGroup.SetExtentTree(extentRootAddr);
		if (status != B_OK)
			return status;
		status = blockGroup.LoadExtent(fTree, false);
		if (status != B_OK)
			return status;
	}

	if (fTree->IsEmpty())	// 4 backup roots is 0
		return B_OK;

	uint64 lowerBound = blockGroup.Start();
	uint64 upperBound = blockGroup.End();
	status = fTree->FillFreeExtents(lowerBound, upperBound);
	if (status != B_OK) {
		ERROR("ExtentAllocator::_LoadExtentTree() could not fill free extents"
			"start %" B_PRIu64 " end %" B_PRIu64 "\n", lowerBound,
			upperBound);
		return status;
	}

	if (fStart > lowerBound)
		fStart = lowerBound;
	if (fEnd < upperBound)
		fEnd = upperBound;
	return B_OK;
}


status_t
ExtentAllocator::Initialize()
{
	TRACE("ExtentAllocator::Initialize()\n");
	if (fTree == NULL)
		return B_NO_MEMORY;

	status_t status = _LoadExtentTree(BTRFS_BLOCKGROUP_FLAG_DATA);
	if (status != B_OK) {
		ERROR("ExtentAllocator:: could not load extent tree (data)\n");
		return status;
	}

	status = _LoadExtentTree(BTRFS_BLOCKGROUP_FLAG_SYSTEM);
	if (status != B_OK) {
		ERROR("ExtentAllocator:: could not load extent tree (system)\n");
		return status;
	}

	status = _LoadExtentTree(BTRFS_BLOCKGROUP_FLAG_METADATA);
	if (status != B_OK) {
		ERROR("ExtentAllocator:: could not load extent tree (metadata)\n");
		return status;
	}

	fTree->DumpInOrder();
	return B_OK;
}


/* Allocate extent that
 * startwith or after "start" and has size >= "size"
 * For now the allocating strategy is "first fit"
 */
status_t
ExtentAllocator::_Allocate(uint64& found, uint64 start, uint64 size,
	uint64 type)
{
	TRACE("ExtentAllocator::_Allocate() start %" B_PRIu64 " size %" B_PRIu64
		" type %" B_PRIu64 "\n", start, size, type);
	CachedExtent* chosen;
	status_t status;
	while (true) {
		status = fTree->FindNext(&chosen, start, size, type);
		if (status != B_OK)
			return status;

		if (TreeDefinition().Compare(start, chosen) == 0) {
			if (chosen->End() - start >= size) {
				// if covered and have enough space for allocating
				break;
			}
			start = chosen->End();	// set new start and retry
		} else
			start = chosen->offset;
	}

	TRACE("ExtentAllocator::_Allocate() found %" B_PRIu64 "\n", start);
	chosen = CachedExtent::Create(start, size,
		type | BTRFS_EXTENT_FLAG_ALLOCATED);
	status = fTree->AddExtent(chosen);
	if (status != B_OK)
		return status;

	found = start;
	return B_OK;
}


// Allocate block for metadata
// flags is BLOCKGROUP's flags
status_t
ExtentAllocator::AllocateTreeBlock(uint64& found, uint64 start, uint64 flags)
{
	// TODO: implement more features here with flags, e.g DUP, RAID, etc

	BlockGroup blockGroup(fVolume->ExtentTree());
	status_t status = blockGroup.Initialize(flags);
	if (status != B_OK)
		return status;
	if (start == (uint64)-1)
		start = blockGroup.Start();

	// normalize inputs
	uint64 remainder = start % fVolume->BlockSize();
	if (remainder != 0)
		start += fVolume->BlockSize() - remainder;

	status = _Allocate(found, start, fVolume->BlockSize(),
		BTRFS_EXTENT_FLAG_TREE_BLOCK);
	if (status != B_OK)
		return status;

	// check here because tree block locate in 2 blockgroups (system and
	// metadata), and there might be a case one can get over the limit.
	if (found >= blockGroup.End())
		return B_BAD_DATA;
	return B_OK;
}


// Allocate block for file data
status_t
ExtentAllocator::AllocateDataBlock(uint64& found, uint64 size, uint64 start,
	uint64 flags)
{
	// TODO: implement more features here with flags, e.g DUP, RAID, etc

	if (start == (uint64)-1) {
		BlockGroup blockGroup(fVolume->ExtentTree());
		status_t status = blockGroup.Initialize(flags);
		if (status != B_OK)
			return status;
		start = blockGroup.Start();
	}

	// normalize inputs
	uint64 remainder = start % fVolume->SectorSize();
	if (remainder != 0)
		start += fVolume->SectorSize() - remainder;
	size = size / fVolume->SectorSize() * fVolume->SectorSize();

	return _Allocate(found, start, size, BTRFS_EXTENT_FLAG_DATA);
}