#include <stdint.h>
#include <stddef.h>
#include <memory.h>

#define MEMORY_INTERNAL
#include <memory/virtalloc.h>

struct addr_node {
	uintptr_t start;
	uintptr_t end;
	struct addr_node *next;
	struct addr_node *prev;
	uint8_t is_alloc;
	uint8_t is_bss;
};

#define BOOTSTRAP_BSS_NODES 16
static uint8_t bss_nodes = 0;
static struct addr_node nodes[BOOTSTRAP_BSS_NODES];

static struct addr_node *start_node;

static struct addr_node *alloc_node(void) {
	struct addr_node *node = NULL;
	if (bss_nodes >= BOOTSTRAP_BSS_NODES) {
		node = kalloc(sizeof(struct addr_node));
		if (node == NULL)
			return NULL;
		node->is_bss = false;
	} else {
		node = &nodes[bss_nodes];
		bss_nodes += 1;
		node->is_bss = true;
	}
	return node;
}

static void free_node(struct addr_node *node) {
	if (!node->is_bss)
		kfree(node);
}

void virtaddr_init(void) {
	struct addr_node init = {
		.start = 0x400000, // third page table
		.end = 0x1000000000000, // 48bit memory address max
		.next = NULL,
		.prev = NULL,
		.is_alloc = false,
		.is_bss = true,
	};
	nodes[0] = init;
	start_node = &nodes[0];
	bss_nodes++;
}

void *virtaddr_alloc(int n_pages) {

	if (n_pages < 1)
		return NULL;

	long n_length = n_pages * PAGE_SIZE;
	struct addr_node *node = start_node;

	for (; node != NULL ; node = node->next) {

		if (node->is_alloc)
			continue;

		long length = node->end - node->start;
		if (length >= n_length) {
			struct addr_node *new = alloc_node();
			if (node == NULL) {
				return NULL;
			}
			if (node->prev != NULL) {
				node->prev->next = new;
			}
			new->next = node;
			node->prev = new;
			new->start = node->start;
			new->end = new->start + n_length;
			node->start = new->end;
			new->is_alloc = true;
			new->next = node;
			return (void *) new->start;
		}
	}

	return NULL;
}

static void merge_back(struct addr_node *node) {
	while(node->prev && !node->is_alloc) {
		struct addr_node *temp = node->prev;
		node->start = node->prev->start;
		node->prev = node->prev->prev;
		free_node(temp);
	}
	if (node->prev == NULL) {
		start_node = node;
	}
}

static void merge_forward(struct addr_node *node) {
	while(node->next && !node->is_alloc) {
		struct addr_node *temp = node->next;
		node->end = node->next->end;
		node->next = node->next->next;
		free_node(temp);
	}
}

long virtaddr_free(void *virtaddr) {

	uintptr_t virt = (uintptr_t) virtaddr;

	if (virt % PAGE_SIZE)
		return -1; // not page aligned, we did not give this out!!!

	struct addr_node *node = start_node;

	for (; node != NULL; node = node->next) {
		if (node->start == virt) {
			int length = node->end - node->start;
			int pages = length / PAGE_SIZE;
			merge_back(node);
			merge_forward(node);
			return pages;
		}
	}

	return -1;
}