#include <lcrash/mm/kmalloc.h>

#include <lcrash/mm/phys.h>

/**
 * A block in the kernel heap. the performance is abysmal, this structure may be larger
 * than a lot of kernel objects so we'll have about half the memory we could have ^-^
 */
struct KernelHeapBlock {
	/// Next block
	struct KernelHeapBlock* Next;
	
	/// Previous block
	struct KernelHeapBlock* Prev;

	/// Next free block
	struct KernelHeapBlock* NextFree;

	/// Previous free block
	struct KernelHeapBlock* PrevFree;

	/// Block size
	u64 Size;

	/// Block is free
	bool Free;
};

/**
 * A structure that describes a specific kernel heap, yes, we can have more than one and,
 * yes, most will do the same thing
 */
struct KernelHeapZone {
	/// Next heap zone
	struct KernelHeapZone* Next;

	/// Previous heap zone
	struct KernelHeapZone* Prev;

	/// First heap frame
	struct KernelHeapFrame* FirstFrame;

	/// Last heap frame
	struct KernelHeapFrame* LastFrame;

	/// First block
	struct KernelHeapBlock* First;

	/// Last block
	struct KernelHeapBlock* Last;

	/// First free block
	struct KernelHeapBlock* FirstFree;

	/// Last free block
	struct KernelHeapBlock* LastFree;
};

/**
 * A block of physical memory owned by a specific heap, where is it stored exactly? In
 * the heap of course!
 */
struct KernelHeapFrame {
	/// The block
	struct PmemAllocation Block;

	/// The block's size
	u64 Size;

	/// Linked list jumpscare
	struct KernelHeapFrame* NextFrame;

	/// I wonder if this is a pointer that points to the previous frame in a linked
	/// list of the structure which stores kernel heap frames
	struct KernelHeapFrame* PrevFrame;
};

struct KernelHeapZone* KernelHeapFirst = 0;
struct KernelHeapZone* KernelHeapLast = 0;

/**
 * \todo Maybe split the zone allocation code out of the thing
 */
int KernelHeapInitialize() {
	// Allocate uh this amount of memory for the heap
	struct PmemAllocation Block;
	if (PmemAllocateBlock(0x100000, &Block)) return -1;

	// Make our zone
	struct KernelHeapZone* Zone = Block.Address;
	Zone->Next = 0;
	Zone->Prev = 0;
	KernelHeapFirst = Zone;
	KernelHeapLast = Zone;

	// Make the frame block
	struct KernelHeapBlock* FrameBlock = Block.Address + sizeof(struct KernelHeapZone);
	FrameBlock->Prev = 0;
	FrameBlock->PrevFree = 0;
	FrameBlock->Size = sizeof(struct KernelHeapFrame);
	FrameBlock->Free = false;

	// Make the frame structure
	struct KernelHeapFrame* Frame = Block.Address 
				      + sizeof(struct KernelHeapZone)
				      + sizeof(struct KernelHeapBlock);
	Frame->Block = Block;
	Frame->Size = 0x100000;
	Frame->NextFrame = 0;
	Frame->PrevFrame = 0;

	// Make the root block
	struct KernelHeapBlock* RootBlock = (struct KernelHeapBlock*)((void*)Frame + sizeof(struct KernelHeapFrame));
	RootBlock->Prev = FrameBlock;
	FrameBlock->Next = RootBlock;
	FrameBlock->NextFree = RootBlock;
	RootBlock->Next = 0;
	RootBlock->PrevFree = 0;
	RootBlock->NextFree = 0;
	RootBlock->Free = true;
	RootBlock->Size = 0x100000
		    - sizeof(struct KernelHeapZone)
		    - sizeof(struct KernelHeapBlock)
		    - sizeof(struct KernelHeapFrame)
		    - sizeof(struct KernelHeapBlock);

	// Fill out the rest of the zone
	Zone->FirstFrame = Frame;
	Zone->LastFrame = Frame;
	Zone->FirstFree = RootBlock;
	Zone->LastFree = RootBlock;
	Zone->First = FrameBlock;
	Zone->Last = RootBlock;

	// Tada!
	return 0;
}

void* KernelHeapAlloc(u32 Size, u32 Flags) {
	// Find a free block
	for (struct KernelHeapZone* Zone = KernelHeapFirst; Zone != 0; Zone = Zone->Next) {
		for (struct KernelHeapBlock* Block = Zone->FirstFree; Block != 0; Block = Block->NextFree) {
			if (Block->Size >= Size) {
				s64 SplitSize = (s64)Block->Size - Size - (sizeof(struct KernelHeapBlock) * 2) - 8;

				if (SplitSize > 0) {
					struct KernelHeapBlock* NewBlock = (struct KernelHeapBlock*)((void*)Block + sizeof(struct KernelHeapBlock) + Size);
					NewBlock->Prev = Block;
					NewBlock->Next = 0;
					NewBlock->PrevFree = Block->PrevFree;
					NewBlock->NextFree = Block->NextFree;
					NewBlock->Free = true;
					NewBlock->Size = (u64)SplitSize;

					// Set zone stuff
					if (Block->NextFree == 0) Zone->LastFree = NewBlock;
					if (Block->PrevFree == 0) Zone->FirstFree = NewBlock;
					if (Block->Next == 0) Zone->Last = NewBlock;

					// Set next free in other blocks
					for (struct KernelHeapBlock* ModBlock = Block->Prev; ModBlock != 0 && (Block->NextFree != 0 || Block->NextFree != Block); ModBlock = ModBlock->Prev) {
						ModBlock->NextFree = NewBlock;
					}

					// Set prev free in other blocks
					for (struct KernelHeapBlock* ModBlock = Block->Next; ModBlock != 0 && (Block->PrevFree != 0 || Block->PrevFree != Block); ModBlock = ModBlock->Next) {
						ModBlock->PrevFree = NewBlock;
					}

					// Block stuff
					Block->Next = NewBlock;
				}

				Block->Free = 0;
				Block->Size = Size;

				return (void*)Block + sizeof(struct KernelHeapBlock);
			}
		}
	}

	return 0;
}

void KernelHeapFree(void* Address) {
	
}