#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];

#define PAGE_SIZE sizeof(uint64_t) * 512

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

struct page_mem {
	void *ptr;
	size_t size;
	struct page_mem *next;
} pagelist = {
	.ptr = &kernel_end,
	.size = 0x100000,
	.next = NULL
};


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

static struct memory *split_block(struct memory *block, size_t size) {
	struct memory *nextblock = (void *)((uintptr_t)block->block + size);
	*nextblock = (struct memory) {
		.size = block->size - size - sizeof(*nextblock),
		.free = true,
		.prev = block,
		.next = block->next
	};
	block->next = nextblock;
	block->size = size;
	block->free = false;
	return nextblock;
}

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)
			break;

		if (block->size < (2 * sizeof(struct memory)) + size + (alignment - offset))
			continue;

		struct memory *new_block = split_block(block, (alignment - offset - sizeof(*block)));
		block = new_block;
		break;
	}

	if (!block) {
		//print_blocks();
		vga_puts("\n\noom");
		while(1);
	}
	
	split_block(block, size);
	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);
	}
	
	split_block(block, size);
	return block->block;
}

void kfree(void *ptr) {
	struct memory *header = (void *)((uintptr_t)ptr - sizeof(struct memory));
	if (header->next && header->next == (void*)((uintptr_t)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*)((uintptr_t)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;
}

static uint64_t alloc_page(void) {
	uint64_t *page;

	if (blocklist) {
		page = kmalloc_a(PAGE_SIZE, 0x1000);
		memset(page, 0, PAGE_SIZE);
		return (uint64_t)page | 0x3;
	}

	if (pagelist.size < PAGE_SIZE) {
		if (!pagelist.next) {
			vga_puts("oom");
			while(1);
		}
		pagelist = *pagelist.next;
	}

	if ((uintptr_t) pagelist.ptr & 0xfff) {
		pagelist.ptr = (void *) (((uintptr_t) pagelist.ptr & ~0xfff) + 0x1000);
		pagelist.size -= (0x1000 - (uintptr_t) pagelist.ptr % 0x1000);
	}

	page = pagelist.ptr;
	pagelist.ptr = (void *) ((uintptr_t) pagelist.ptr + PAGE_SIZE);
	pagelist.size -= PAGE_SIZE;

	memset(page, 0, PAGE_SIZE);
	return (uint64_t)page | 0x3;
}

static uint64_t *addr_to_page(uint64_t addr) {
	uint64_t *page = pml4;
	size_t offset = 39;
	size_t index = (addr >> offset) & 0x1ff;

	for (size_t i = 0; i < 3; i++) {
		if (!(page[index] & 0x1))
			page[index] = alloc_page();
		page = (uint64_t*)(page[index] & ~0xfff);
		index = (addr >> (offset -= 9)) & 0x1ff;
	}
	return &page[index];
}

void memory_map(uintptr_t addr, size_t size) {
	for (uintptr_t ptr = addr & ~0xfff; ptr < addr + size; ptr += 0x1000) {
		*addr_to_page(ptr) = ptr | 0x3 | (1 << 6);
	}
}

bool paging_init(void) {
	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];

		vga_puts("region addr ");
		vga_putx(mmap->addr);
		vga_puts(" len ");
		vga_putx(mmap->len);
		switch (mmap->type) {
			case MULTIBOOT_MEMORY_AVAILABLE:
				vga_puts(" avail");
				break;
			case MULTIBOOT_MEMORY_NVS:
				vga_puts(" nvs");
				break;
			case MULTIBOOT_MEMORY_BADRAM:
				vga_puts(" badram");
				break;
			case MULTIBOOT_MEMORY_RESERVED:
				vga_puts(" reserv");
				break;
			case MULTIBOOT_MEMORY_ACPI_RECLAIMABLE:
				vga_puts(" acpi");
				break;
		}
		vga_putc('\n');

		if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE && mmap->addr >= 0x100000
				&& mmap->addr + mmap->len > (uintptr_t)&kernel_end + 0x100000) {
			struct page_mem *next_page = mmap->addr > (uintptr_t)&kernel_end + 0x100000 ?
				(void*) mmap->addr : (void *) ((uintptr_t)&kernel_end + 0x100000);

			*next_page = (struct page_mem) {
				.ptr = next_page,
				.size = mmap->len,
				.next = NULL
			};

			struct page_mem *last = &pagelist;
			for (; last->next; last = last->next);
			last->next = next_page;
		}

		for (size_t addr = mmap->addr & ~0xfff; addr <= mmap->addr + mmap->len; addr += 0x1000)
			*addr_to_page(addr) = addr | 0x3;
	}

	for (struct page_mem *region = &pagelist; region; region = pagelist.next) {
		pagelist = *region;
		struct memory *mem_block = pagelist.ptr;
		(*mem_block) = (struct memory) {
			.size = pagelist.size - sizeof(struct memory),
			.free = true
		};

		if (!blocklist) {
			blocklist = mem_block;
			continue;
		}

		struct memory *free_block = blocklist;
		while (free_block->next)
			free_block = free_block->next;
		free_block->next = mem_block;
		mem_block->prev = free_block;
	}

	return true;
}