/*
 * This file controls the page frame allocator.  It is the core of the physical
 * memory manager.  The kernel heap and paging systems both sit on top of this.
 * It keeps track of which frames of physical memory a free and hands them out
 * as needed.
 */

#include <stddef.h>
#include <stdint.h>
#include <sys/types.h>
#include <nucleus/kernel.h>
#include <nucleus/lib.h>
#include <nucleus/memory.h>

#define INDEX(a) ((a)/32)
#define OFFSET(a) ((a)%32)

/* Types of an E820 Memory Map Entry */
enum E820Type {
	E820_NONE,
	E820_USABLE,
	E820_RESERVED,
	E820_RECLAIMABLE,
	E820_NVS,
	E820_BAD,
};

/* Structure of an E820 Memory Map Entry */
struct E820Entry {
	uint32_t size;
	uint32_t base, baseHigh;
	uint32_t length, lengthHigh;
	uint32_t type;
} __attribute__((packed));

/* Descriptor of a Frame Region */
struct FrameRegion {
	uintptr_t base;
	size_t numFrames, usedFrames;
	struct FrameRegion *next;
	uint32_t bitmap[];
};

size_t numFrames, usedFrames;
struct FrameRegion *regions;

/* Count the number of pages a number of bytes occupies */
static size_t
page_count(size_t bytes)
{
	size_t pages = bytes / PAGE_SIZE;
	if (bytes & (PAGE_SIZE - 1))
		return pages + 1;
	return pages;
}

/* Set a bit in a frame's bitset */
static void
set_frame(struct FrameRegion *region, off_t index)
{
	region->bitmap[INDEX(index)] |= (1 << OFFSET(index));
	region->usedFrames++;
	usedFrames++;
}

/* Clear a bit in a frame's bitset */
static void
clear_frame(struct FrameRegion *region, off_t index)
{
	region->bitmap[INDEX(index)] &= ~(1 << OFFSET(index));
	region->usedFrames--;
	usedFrames--;
}

/* Test a bit in a frame's bitset */
static uint32_t
test_frame(struct FrameRegion *region, off_t index)
{
	return (region->bitmap[INDEX(index)] & (1 << OFFSET(index)));
}

/* Find first free frame */
static off_t
find_frame(struct FrameRegion *region)
{
	off_t i, j;
	for (i = 0; i < INDEX(region->numFrames); i++) {
		if (!~region->bitmap[i])
			continue;
		for (j = 0; j < 32; j++) {
			if ((region->bitmap[i] >> j) & 1)
				continue;
			return (i * 32) + j;
		}
	}
	return (off_t) -1;
}

/* Allocate a page frame */
uintptr_t
alloc_frame(void)
{
	int idx;
	/* Walk the regions, first fit */
	struct FrameRegion *region;
	for (region = regions; region; region = region->next) {
		idx = find_frame(region);
		if (idx != -1)
			break;
	}
	if (idx == -1) {
		panic("Could not allocate frame!");
		return 0x00000000;
	}

	set_frame(region, idx);
	return region->base + (idx << 12);
}

/* Free a page frame */
void
free_frame(uintptr_t frame)
{
	/* Walk the regions */
	struct FrameRegion *region;
	for (region = regions; region->next; region = region->next) {
		if (region->next->base >= frame)
			break;
	}
	clear_frame(region, (frame - region->base) >> 12);
}

/* Setup the page frame allocator */
void
init_frames(uint32_t size, void *addr)
{
	/*
	 * When the kernel starts, the bootloader passes the kernel a map of
	 * memory.  This map must be read so the memory manager can avoid bad
	 * areas of memory and areas that are mapped to hardware, ACPI, etc.
	 */
	int i, j;
	struct E820Entry *memMap = addr;
	struct FrameRegion *prev = NULL, *head = regions;
	uintptr_t bumpAlloc = PAGE_SIZE * 2; /* Leave room for AP trampoline */
	for (i = 0; i < size / sizeof(struct E820Entry); i++) {
		if (memMap[i].baseHigh > 0
		 || memMap[i].base > 0xFFFFFFFF
		 || memMap[i].type != E820_USABLE)
			continue;

		/* Usable region - create bitmap */
		size_t frameSize = memMap[i].length / PAGE_SIZE;
		size_t bitmapSize = (frameSize / 8);
		head = (struct FrameRegion *) bumpAlloc;
		bumpAlloc += bitmapSize + sizeof(struct FrameRegion);
		head->base = memMap[i].base;
		if (sizeof(head->base) > sizeof(memMap[i].base))
			head->base += memMap[i].baseHigh;
		head->numFrames = frameSize;
		head->usedFrames = 0;

		memset(head->bitmap, 0, bitmapSize);
		/* Bits that don't refer to real frames must be set */
		for (j = OFFSET(head->numFrames); j < 32; j++)
			set_frame(head, head->numFrames + j);

		if (prev) prev->next = head;
		else regions = head;
		prev = head;
	}

	/* Regions to be remapped */
	struct {uint32_t start, end;} remaps[] = {
		{.start = 0x0000,   .end = bumpAlloc}, /* PMM bitmaps */
		{.start = 0x100000, .end = 0x180000 }, /* Kernel */
		{.start = 0x200000, .end = 0x800000 }, /* Kernel heap */
	};

	/* Check bitmaps */
	numFrames = usedFrames = 0;
	struct FrameRegion *region;
	for (region = regions; region; region = region->next) {
		numFrames += region->numFrames;
		uintptr_t end = region->base + (region->numFrames * PAGE_SIZE);
		/* Iterate the remaps[] to find overlapping regions */
		for (i = 0; i < sizeof(remaps)/sizeof(remaps[0]); i++) {
			for (j = remaps[i].start;
			     j < remaps[i].end && j >= region->base && j < end;
			     j += PAGE_SIZE)
				set_frame(region, (j - region->base) >> 12);
		}
	}
}