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

#include <stdint.h>

uint32_t *page_directory = &boot_page_directory;
uint32_t page_table_default[1024] __attribute__((aligned(PAGE_SIZE)));
uint32_t multiboot_page_table[1024] __attribute__((aligned(PAGE_SIZE)));
uint32_t frame_zones_page_table[1024] __attribute__((aligned(PAGE_SIZE)));
uint32_t mem_size;
multiboot_memory_map_t *mmap_addr;
multiboot_uint32_t mmap_length;
struct frame_zone *head;

static void init_multiboot(multiboot_info_t *mbd, uint32_t magic)
{
	if (magic != MULTIBOOT_BOOTLOADER_MAGIC)
		kpanic("invalid magic number! (git good skill issue)");

	init_page_table(multiboot_page_table, 0);
	page_directory[1023] =
	    ((uint32_t)multiboot_page_table - HEAP_END) | 0x03;
	const size_t mbd_size = CEIL(
	    (uint32_t)mbd % PAGE_SIZE + sizeof(multiboot_info_t), PAGE_SIZE);

	// Index multiboot_info_t struct
	for (uint32_t i = 0; i < mbd_size; i++)
		multiboot_page_table[i] =
		    (((uint32_t)mbd + PAGE_SIZE * i) & PAGE_MASK) | INIT_FLAGS;
	multiboot_info_t *mbd_virt =
	    (multiboot_info_t *)(GET_PAGE_ADDR(1023, 0) +
	                         (uint32_t)mbd % PAGE_SIZE);

	// Index mbd->mmap_addr pointers
	for (uint32_t i = 0; i < mbd_virt->mmap_length; i++)
		multiboot_page_table[i + mbd_size] =
		    ((mbd_virt->mmap_addr + i * PAGE_SIZE) & PAGE_MASK) |
		    INIT_FLAGS;
	mmap_addr = (multiboot_memory_map_t *)(GET_PAGE_ADDR(1023, mbd_size) +
	                                       (uint32_t)mbd_virt->mmap_addr %
	                                           PAGE_SIZE);
	mmap_length = mbd_virt->mmap_length / sizeof(multiboot_memory_map_t);
}

static void add_frame_node(multiboot_memory_map_t *mmmt)
{
	static uint32_t index;
	void *zone = (void *)mmmt->addr;

	init_page_table(frame_zones_page_table, 0);
	page_directory[1022] =
	    ((uint32_t)frame_zones_page_table - HEAP_END) | 0x03;
	frame_zones_page_table[index] =
	    ((uint32_t)zone & PAGE_MASK) | INIT_FLAGS;

	struct frame_zone *current =
	    (struct frame_zone *)GET_PAGE_ADDR(1022, index++);
	memset(current, 0, sizeof(struct frame_zone));
	current->addr = (void *)mmmt->addr;
	current->frame_table = (uint32_t *)current + sizeof(struct frame_zone);

	const size_t frame_table_size =
	    CEIL(current->size - sizeof(struct frame_zone), PAGE_SIZE * 32);
	for (uint32_t i = index; index - i < CEIL(frame_table_size, PAGE_SIZE);
	     index++)
		frame_zones_page_table[index] =
		    ((uint32_t)zone + PAGE_SIZE & PAGE_MASK) | INIT_FLAGS;
	// glhf reading this bozo
	current->len = mmmt->len;
	current->size = (mmmt->len - (sizeof(struct frame_zone) +
	                              (mmmt->len / PAGE_SIZE) / 32));
	current->remaining_frames = current->size / PAGE_SIZE;
	current->first_free_frame = 0;
	current->next = NULL;

	if (!head) {
		head = current;
		return;
	}
	struct frame_zone *it = head;
	while (it->next)
		it = it->next;
	it->next = current;
}

static void init_frame_zones(void)
{
	for (uint32_t i = 0; i < mmap_length; i++) {
		multiboot_memory_map_t *mmmt =
		    (multiboot_memory_map_t *)mmap_addr + i;
		if (mmmt->type == MULTIBOOT_MEMORY_AVAILABLE)
			add_frame_node(mmmt);
	}
}

void init_memory(multiboot_info_t *mbd, uint32_t magic)
{
	assert(page_directory);
	for (uint16_t i = 0; i < 0x300; i++)
		page_directory[i] = 0x02;
	init_page_table(page_table_default, 0);
	page_directory[0] = ((uint32_t)page_table_default - HEAP_END) | 0x03;
	init_multiboot(mbd, magic);
	init_frame_zones();
}