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

#include "kprintf.h"
#include "memory.h"
#include "utils.h"

#define MAX_TLB_ENTRIES 32

static uint32_t *current_page_table;
static uint16_t current_pd_index;

static int16_t find_next_block(size_t nb_pages)
{
	for (uint16_t pd_index = 0; pd_index < 768; pd_index++) {
		if (page_directory[pd_index] == 0x02) {
			if (create_page_table(pd_index) < 0)
				return -2;
		}
		current_pd_index = pd_index;
		current_page_table =
		    (uint32_t *)((PT_START + pd_index) * PAGE_SIZE);
		for (uint16_t i = 0; i + nb_pages < PT_SIZE; i++) {
			uint16_t j;
			for (j = 0; current_page_table[i + j] >> 12 == i + j &&
			            j < nb_pages;
			     j++)
				;
			if (j == nb_pages)
				return i;
			i += j;
		}
	}
	return -1;
}

void *alloc_pages(size_t size)
{
	const uint32_t nb_pages = CEIL(size, PAGE_SIZE);
	const int16_t index = find_next_block(nb_pages);

	if (index < 0) {
		kprintf(KERN_CRIT "%d: Not enough pages (max: %d)\n", index,
		        PT_SIZE);
		return NULL;
	}
	for (size_t i = index; i - (size_t)index < nb_pages; i++) {
		void *frame = alloc_frames(PAGE_SIZE);
		if (!frame) {
			for (size_t j = index; j < i; j++)
				free_frames(
				    (void *)(current_page_table[j] >> 12),
				    PAGE_SIZE);
			return NULL;
		}
		current_page_table[i] =
		    ((uint32_t)frame & PAGE_MASK) | INIT_FLAGS;
	}
	return (void *)(((current_pd_index * 1024) + index) * PAGE_SIZE);
}

int free_pages(void *page_ptr, size_t size)
{
	const uint32_t nb_pages = CEIL(size, PAGE_SIZE);
	const uint32_t page_index = (uint32_t)page_ptr / PAGE_SIZE;

	if ((uint32_t)page_ptr > PT_SIZE * PAGE_SIZE) {
		kprintf(KERN_WARNING "Address out of range\n");
		return -1;
	} else if ((uint32_t)page_ptr % PAGE_SIZE) {
		kprintf(KERN_WARNING "Invalid address\n");
		return -1;
	} else if (page_index + nb_pages > PT_SIZE) {
		kprintf(KERN_WARNING "Invalid number of frames\n");
		return -1;
	}
	for (size_t i = page_index; i < page_index + nb_pages; i++) {
		free_frames((void *)(current_page_table[i] >> 12), PAGE_SIZE);
		current_page_table[i] = i << 12 | INIT_FLAGS;
	}
	return 0;
}