#include <mm/paging.h>
#include <mm/kernel_heap.h>
#include <toxic/errno.h>
#include <toxic/config.h>
#include <toxic/vga.h>

void paging_load_directory(uint32_t *directory);

static struct page_directory *current_directory;

struct page_directory*
paging_new_directory(uint8_t flags)
{
	int i, j;
	int offset = 0;

	/* Allocate the global directory */
	struct page_directory *new_directory = kzalloc(sizeof(struct page_directory));
	uint32_t *directory = kzalloc(sizeof(uint32_t) * PAGING_ENTRIES_PER_TABLE);

	/* Allocate a new page table for each page directory entry */
	for (i = 0; i < PAGING_ENTRIES_PER_TABLE; i++) {
		uint32_t *page_table = kzalloc(sizeof(uint32_t) * PAGING_ENTRIES_PER_TABLE);

		/*
		 * Each PTE is initialized now with the linear physical address.
		 * I don't think this is a good idea, but I'll change it later.
		 *
		 * In theory we should not allocate memory for PTEs not used, so
		 * this is in general a waste of space. A page table should be allocated only
		 * when we actually need to allocate space there.
		 *
		 * Also, PTEs should point to NULL instead of defaulting to a
		 * physical address.
		 */
		for (j = 0; j < PAGING_ENTRIES_PER_TABLE; j++)
			page_table[j] = (offset + (j * PAGE_SIZE)) | flags;

		offset += (PAGING_ENTRIES_PER_TABLE * PAGE_SIZE);

		/*
		 * Set the whole directory as WRITABLE so we don't
		 * prevent any page to be writable
		 */
		directory[i] = (uint32_t)page_table | flags | PAGING_IS_WRITABLE;
	}

	new_directory->directory = directory;
	return new_directory;
}

bool
paging_is_aligned(void *addr)
{
	return ((uint32_t)addr % PAGE_SIZE) == 0;
}

uint32_t *
paging_get_directory(struct page_directory *pd)
{
	return pd->directory;
}

void
paging_switch(struct page_directory *directory)
{
	paging_load_directory(paging_get_directory(directory));
	current_directory = directory;
}

int
paging_get_indexes(
	void *vaddr,
	uint32_t *dir_idx,
	uint32_t *table_idx)
{
	int ret = 0;

	if (!paging_is_aligned(vaddr)) {
		ret = -EINVAL;
		goto out;
	}

	*dir_idx = (uint32_t)vaddr >> PAGING_DIRECTORY_BITS;
	*table_idx = ((uint32_t)vaddr >> PAGING_TABLE_BITS) & PAGING_TABLE_MASK;

out:
	return ret;
}

/*
 * Configure a PTE to map virtual address vaddr to physical address paddr
 * setting the entry bits to flags
 */
int
paging_map_vaddr(
	struct page_directory *page_dir,
	void *vaddr,
	void *paddr,
	uint32_t flags)
{
	uint32_t *directory;
	uint32_t *page_table;
	uint32_t dir_idx = 0;
	uint32_t table_idx;
	uint32_t dir_entry;
	uint32_t pte = (uint32_t)paddr;
	int ret = 0;

	if (!paging_is_aligned(vaddr) || !paging_is_aligned(paddr)) {
		return -EINVAL;
	}

	ret = paging_get_indexes(vaddr, &dir_idx, &table_idx);
	if (ret < 0)
		return ret;

	directory = paging_get_directory(page_dir);
	dir_entry = (uint32_t)directory[dir_idx];

	page_table = (uint32_t *)(dir_entry & (~PAGING_PAGE_OFFSET_MASK));
	page_table[table_idx] = pte | flags;

	return ret;
}