#include <memory-add.h>
#include <debugio.h>
#include <stdlib.h>

#define MEMHEAP_MAGIC 0x50DBE514
#define MEMHEAP_INDEX_SIZE 0x20000
#define MEMHEAP_MINIM_SIZE 0x70000

typedef struct
{
	uint32 Magic;
	// bit 0: used  bit 1: reserved
	uint8 Used;
	uint32 Size;
} MemHeapHeader;

typedef struct
{
	uint32 Magic;
	MemHeapHeader* Header;
} MemHeapFooter;

MemHeap* KernelHeap;

uint32 MemHeapFindSmallestHole (uint32 size, uint8 page_align, MemHeap* heap)
{
	uint32 i;
	for (i = 0; i < heap->Index.Size; i++)
	{
		MemHeapHeader* head = (MemHeapHeader*) OrderedArrayLookup(i, &heap->Index);

		if (page_align)
		{
			uint32 location = (uint32)(head) + sizeof(MemHeapHeader);

			// page align it
			uint32 offset = 0x1000 - (location & 0xfff);
			offset &= 0xfff;

			if (head->Size - offset >= size) return i;
		}

		else if (head->Size >= size) return i;
	}

	return 0xffffffff;
}

int32 MemHeapCompare (uint32 a, uint32 b)
{
	MemHeapHeader *ha = (MemHeapHeader*)a, *hb = (MemHeapHeader*)b;

	if (ha->Size > hb->Size) return 1;
	else if (ha->Size == hb->Size) return 0;
	return -1;
}


inline void MemHeapHeaderSetup (uint32 size, uint8 used, MemHeapHeader* head)
{
	head->Magic = MEMHEAP_MAGIC;
	head->Used = used;
	head->Size = size;
}

inline void MemHeapFooterSetup (MemHeapHeader* head, MemHeapFooter* foot)
{
	foot->Magic = MEMHEAP_MAGIC;
	foot->Header = head;
}


MemHeap* MemHeapCreate(uint32 start, uint32 end, uint32 max, uint8 flags)
{
	if ((start & 0xfff) || (end & 0xfff)) return NULL;

	MemHeap* heap = (MemHeap*) kmalloc(sizeof(MemHeap));
	heap->Index = OrderedArrayPlace(start, MEMHEAP_INDEX_SIZE, MemHeapCompare);

	start += sizeof(uint32) * MEMHEAP_INDEX_SIZE;

	if (start & 0xfff) start = (start & 0xFFFFF000) + 0x1000;

	heap->StartAddress = start;
	heap->EndAddress = end;
	heap->Flags = flags;
	heap->MaxAddress = max;

	// One large hole
	MemHeapHeader* hole = (MemHeapHeader*)start;
	MemHeapHeaderSetup(end-start, 0, hole);
	OrderedArrayInsert(start, &heap->Index);

	return heap;
}


uint32 MemHeapExpand(uint32 newsz, MemHeap* heap, PageDirectory* pd)
{
	if (newsz <= heap->EndAddress - heap->StartAddress) return heap->EndAddress - heap->StartAddress;

	if (newsz & 0xfff) newsz = (newsz & 0xfffff000) + 0x1000;
	if (newsz + heap->StartAddress >= heap->MaxAddress) return heap->EndAddress - heap->StartAddress;

	uint32 i;
	for (i = heap->EndAddress; i < heap->StartAddress + newsz; i+=0x1000)
		PagingMapPage(MemPhAllocateBlock(), i, heap->Flags, pd);

	heap->EndAddress = heap->StartAddress + newsz;
	return newsz;
}

uint32 MemHeapContract(uint32 newsz, MemHeap* heap, PageDirectory* pd)
{
	if (newsz >= heap->EndAddress - heap->StartAddress) return heap->EndAddress - heap->StartAddress;

	if (newsz & 0xfff) newsz = (newsz & 0xfffff000) + 0x1000;		// page align
	newsz = Max(newsz, MEMHEAP_MINIM_SIZE);

	uint32 i;
	for (i = heap->EndAddress - 0x1000; i > heap->StartAddress + newsz; i-=0x1000)
		PagingUnmapPage(i, pd);

	heap->EndAddress = heap->StartAddress + newsz;
	return newsz;
}


uint32 MemHeapAlloc (uint32 size, uint8 isPageAligned, MemHeap* heap, PageDirectory* pd)
{
	// Sanity check
	if (!size || !heap) return 0;

	// Find a good hole
	uint32 newsize = size + sizeof(MemHeapHeader) + sizeof(MemHeapFooter);
	uint32 i = MemHeapFindSmallestHole(newsize, isPageAligned, heap);

	// Didn't find? Expand heap
	if (i == 0xffffffff)
	{
		uint32 oldLen = heap->EndAddress - heap->StartAddress;
		uint32 oldEnd = heap->EndAddress;
		MemHeapExpand(oldLen + newsize, heap, pd);

		uint32 newLen = heap->EndAddress - heap->StartAddress;

		// Find the last header
		uint32 i = 0, index = 0, addr = 0;
		for (; i < heap->Index.Size; i++)
		{
			uint32 tmp = OrderedArrayLookup(i, &heap->Index);
			if (tmp > addr) {
				addr = tmp; index = i;
			}
		}


		MemHeapHeader* head; MemHeapFooter* foot;

		if (heap->Index.Size == 0)	// No headers? Add one
		{
			head = (MemHeapHeader*)oldEnd;
			MemHeapHeaderSetup(newLen - oldLen, 0, head);
			OrderedArrayInsert(oldEnd, &heap->Index);
		}

		else {						// Modify last header
			head = (MemHeapHeader*) OrderedArrayLookup(index, &heap->Index);
			head->Size = newLen - oldLen;
		}

		foot = (MemHeapFooter*)((uint32)head + head->Size - sizeof(MemHeapFooter));
		MemHeapFooterSetup(head, foot);

		// Try again
		return MemHeapAlloc(size, isPageAligned, heap, pd);
	}

	// Get info about hole
	uint32 origAddr = OrderedArrayLookup(i, &heap->Index);
	MemHeapHeader* origHead =(MemHeapHeader*) origAddr;
	uint32 origSize = origHead->Size;

	// Should split in two?
	if (origSize - newsize <= sizeof(MemHeapHeader) + sizeof(MemHeapFooter))
	{
		size += origSize - newsize;
		newsize = origSize;
	}

	// Should be page aligned
	if (isPageAligned && (origAddr&0xfffff000))
	{
		uint32 newAddr = origAddr + 0x1000 - (origAddr&0xfff) - sizeof(MemHeapHeader);
		MemHeapHeaderSetup(0x1000 - (origAddr&0xfff) - sizeof(MemHeapHeader), 0, origHead);

		MemHeapFooter* holeFooter = (MemHeapFooter*)(newAddr - sizeof(MemHeapFooter));
		MemHeapFooterSetup(origHead, holeFooter);

		origAddr = newAddr;
		origSize -= origHead->Size;
	}

	else OrderedArrayDeleteIndex(i, &heap->Index);

	// Create the new block
	MemHeapHeader* bHead = (MemHeapHeader*)origAddr;
	MemHeapFooter* bFoot = (MemHeapFooter*)(origAddr + sizeof(MemHeapHeader) + size);
	MemHeapHeaderSetup(newsize, 1, bHead);
	MemHeapFooterSetup(bHead, bFoot);

	// Create a new hole at the end of current block if necessary
	if (origSize - newsize > 0)
	{
		// header
		MemHeapHeader* hHead = (MemHeapHeader*)(origAddr + sizeof(MemHeapHeader) +
				sizeof(MemHeapFooter) + size);
		MemHeapHeaderSetup(origSize - newsize, 0, hHead);

		// footer
		MemHeapFooter* hFoot = (MemHeapFooter*)((uint32)hHead + hHead->Size - sizeof(MemHeapFooter));
		if ((uint32)hFoot < heap->EndAddress) MemHeapFooterSetup(hHead, hFoot);

		// add it to index
		OrderedArrayInsert((uint32)hHead, &heap->Index);
	}

	return (uint32)bHead + sizeof(MemHeapHeader);
}

void MemHeapFree (uint32 address, MemHeap* heap, PageDirectory* pd)
{
	// Sanity check
	if (!address || !heap) return;

	MemHeapHeader* head = (MemHeapHeader*) (address - sizeof(MemHeapHeader));
	MemHeapFooter* foot = (MemHeapFooter*) ((uint32)head + head->Size - sizeof(MemHeapFooter));

	if (head->Magic != MEMHEAP_MAGIC) return;
	if (foot->Magic != MEMHEAP_MAGIC) return;

	// Clear used flag
	head->Used = 0;

	uint8 AddToIndex = 1;

	// Unify left (if possible)
	MemHeapFooter* testFoot = (MemHeapFooter*) ((uint32)head - sizeof(MemHeapFooter));
	if (testFoot->Magic == MEMHEAP_MAGIC && testFoot->Header->Used == 0)
	{
		uint32 temp = head->Size;
		head = testFoot->Header;
		foot->Header = head;
		head->Size += temp;

		AddToIndex = 0;
	}

	// Unify right (if possible)
	MemHeapHeader* testHead = (MemHeapHeader*) ((uint32)foot + sizeof(MemHeapFooter));
	if ((uint32)testHead < heap->EndAddress && testHead->Magic == MEMHEAP_MAGIC && testHead->Used == 0)
	{
		head->Size += testHead->Size;
		testFoot = (MemHeapFooter*) ((uint32)testHead + testHead->Size - sizeof(MemHeapFooter));
		foot = testFoot;

		uint32 i;
		for (i = 0; i < heap->Index.Size &&
			(uint32)testHead != OrderedArrayLookup(i, &heap->Index); i++) ;

		if (i < heap->Index.Size) OrderedArrayDeleteIndex(i, &heap->Index);
	}

	// Contract heap (if possible)
	if ((uint32)foot + sizeof(MemHeapFooter) == heap->EndAddress)
	{
		uint32 old = heap->EndAddress - heap->StartAddress;
		uint32 new = MemHeapContract((uint32)head + head->Size - sizeof(MemHeapFooter), heap,pd);

		// Last block still existent
		if (head->Size > (old - new))
		{
			head->Size -= old - new;
			foot = (MemHeapFooter*) ((uint32)head + head->Size - sizeof(MemHeapFooter));
			MemHeapFooterSetup(head, foot);
		}

		// Nope, not existent
		else
		{
			uint32 i;
			for (i=0; i<heap->Index.Size && (uint32)testHead != OrderedArrayLookup(i,&heap->Index); i++) ;

			if (i < heap->Index.Size) OrderedArrayDeleteIndex(i,&heap->Index);
		}
	}

	// Lastly, insert hole in array
	if (AddToIndex) OrderedArrayInsert((uint32)head, &heap->Index);
}