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

#include "assert.h"
#include "kprintf.h"
#include "memory.h"
#include "process.h"
#include "string.h"
#include "utils.h"

static int8_t alloc_pagetable(uint16_t pd_index)
{
	uint32_t *pt = alloc_frame();
	if (!pt)
		return -1;
	PD[pd_index] = (uint32_t)pt | INIT_FLAGS;
	bzero(PTE2VA(1023, pd_index), PAGE_SIZE);
	return 0;
}

static uint32_t *find_next_block(size_t nb_pages)
{
	size_t count = 0;
	for (size_t i = USER_PT_START; i < USER_PT_END; i++) {
		if (!PD[i])
			if (alloc_pagetable(i) < 0)
				return NULL;
		for (size_t j = 0; j < 1024; j++) {
			if (*GET_PTE(i, j)) {
				count = 0;
				continue;
			}
			count++;
			if (count == nb_pages) {
				return GET_PTE(i, j) - (count - 1);
			}
		}
	}
	return NULL;
}

void *valloc_pages(size_t nb_pages)
{
	uint32_t *start = find_next_block(nb_pages);
	if (!start)
		return NULL;
	for (uint32_t i = 0; i < nb_pages; i++) {
		void *frame = alloc_frame();
		if (!frame) {
			for (uint32_t j = 0; j < i; j++)
				free_frame((void *)(((uint32_t)(start + j)) &
				                    PAGE_MASK));
			return NULL;
		}
		assert((uint32_t)frame & PAGE_MASK);
		start[i] = (uint32_t)frame | INIT_FLAGS;
	}
	uint32_t page_index = start - PTE2VA(1023, 0);
	return PTE2VA(page_index / 1024, page_index % 1024);
}

int vfree_pages(void *page_ptr, size_t nb_pages)
{
	const uint32_t page_addr = (uint32_t)page_ptr;
	if (page_addr % PAGE_SIZE) {
		kprintf(KERN_WARNING "Invalid address\n");
		return -1;
	}

	for (uint32_t *pte = VA2PTE(page_addr);
	     pte < VA2PTE(page_addr) + nb_pages; pte++) {
		free_frame((void *)(*pte & PAGE_MASK));
		*pte = 0;
	}
	return 0;
}