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

#include <vga.h>
#include <multiboot.h>
#include <memory.h>
#include <mem.h>

extern uint64_t multiboot_magic;
extern struct multiboot_info *multiboot_header;
extern void *kernel_end;

extern uint64_t pml4[512];

struct memory {
	struct memory *prev, *next;
	bool free;
	size_t size;
	uint8_t block[];
} *blocklist;

void print_blocks() {
	struct memory *block = NULL;
	for (block = blocklist; block; block = block->next) {
		vga_puts("block at: ");
		vga_putx((size_t)block);
		vga_puts(", size: ");
		vga_putx(block->size);
		if (block->free)
			vga_puts(", free");
		vga_putc('\n');
	}
}

void *kmalloc_a(size_t size, size_t alignment) {
	struct memory *block = NULL;
	for (block = blocklist; block; block = block->next) {
		if (!block->free || block->size < size + sizeof(struct memory))
			continue;
		size_t offset = (uintptr_t)block->block % alignment;
		if (offset) {
			if (block->size < (2 * sizeof(struct memory)) + size + (alignment - offset)) {
				continue;
			}

			struct memory *new_block = (void *)block + (alignment - offset);
			*new_block = (struct memory) {
				.size = block->size - (alignment - offset),
				.free = true,
				.prev = block,
				.next = block->next,
			};
			block->next = new_block;
			block->size = alignment - offset;

			block = new_block;
			break;
		}
	}
	if (!block) {
		print_blocks();
		vga_puts("\n\noom");
		while(1);
	}
	
	struct memory *nextblock = (void *)block->block + size;
	*nextblock = (struct memory) {
		.size = block->size - size,
		.free = true,
		.prev = block,
		.next = block->next
	};
	block->next = nextblock;
	block->size = size;
	block->free = false;
	return block->block;
}

void *kmalloc(size_t size) {
	struct memory *block = NULL;
	for (block = blocklist; block && (!block->free || block->size < size + sizeof(struct memory)); block = block->next);
	if (!block) {
		print_blocks();
		vga_puts("\n\noom");
		while(1);
	}
	
	struct memory *nextblock = (void *)block->block + size;
	*nextblock = (struct memory) {
		.size = block->size - size,
		.free = true,
		.prev = block,
		.next = block->next
	};
	block->next = nextblock;
	block->size = size;
	block->free = false;

	return block->block;
}

void kfree(void *ptr) {
	struct memory *header = ptr - sizeof(struct memory);
	if (header->next && header->next == (void*)header->block + header->size && header->next->free) {
		header->size += header->next->size + sizeof(struct memory);
		header->next = header->next->next;
		if (header->next)
			header->next->prev = header;
	}
	if (header->prev && header == (void*)header->prev->block + header->prev->size && header->prev->free) {
		header->prev->size += header->size + sizeof(struct memory);
		header->prev->next = header->next;
		if (header->next)
			header->next->prev = header->prev;
		header = header->prev;
	}

	header->free = true;
}

uint64_t alloc_page() {
	uint64_t *page = kmalloc_a(sizeof(uint64_t) * 512, 0x1000);
	memset(page, 0, sizeof(uint64_t) * 512);
	return (uint64_t)page | 0x3;
}

uint64_t *addr_to_page(uint64_t addr) {
	size_t index = (addr >> 39) & 0x1ff;
	if (!(pml4[index] & 0x1))
		pml4[index] = alloc_page();

	uint64_t *pdpr = (uint64_t*)(pml4[index] & ~0xfff);
	index = (addr >> 30) & 0x1ff;
	if (!(pdpr[index] & 0x1))
		pdpr[index] = alloc_page();

	uint64_t *pd = (uint64_t*)(pdpr[index] & ~0xfff);
	index = (addr >> 21) & 0x1ff;
	if (!(pd[index] & 0x1))
		pd[index] = alloc_page();

	uint64_t *pt = (uint64_t*)(pd[index] & ~0xfff);
	index = (addr >> 12) & 0x1ff;
	return &pt[index];
};

bool paging_init() {
	blocklist = (void *)0x7e00;
	*blocklist = (struct memory) {
		.size = 0x7ffff - 0x7e00,
		.free = true
	};

	if (multiboot_magic != 0x2badb002)
		return false;

	for (size_t i = 0; i < multiboot_header->mmap_length / sizeof(struct multiboot_mmap_entry); i++) {
		struct multiboot_mmap_entry *mmap = &((struct multiboot_mmap_entry *) (uintptr_t) multiboot_header->mmap_addr)[i];

		if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE && mmap->addr > 0x1000) {
			struct memory *header = mmap->addr > (uint64_t)&kernel_end ? (void*) mmap->addr : (void *)&kernel_end;

			// map the first header, in case we need to use this region for paging itself
			uint64_t *page = addr_to_page((uintptr_t) header);
			*page = ((uintptr_t)header & ~0xfff) | 0x3;

			*header = (struct memory) {
				.size = mmap->len,
				.free = true
			};
			struct memory *last;
			for (last = blocklist; last->next; last = last->next);
			last->next = header;
			header->prev = last;
		}

		size_t end = mmap->addr + mmap->len;
		for (size_t addr = mmap->addr & ~0xfff; addr < end; addr += 0x1000) {
			uint64_t *page = addr_to_page(addr);
			*page = addr | 0x3;
		}
	}
	return true;
}