#include <alloc.h>
#include <paging.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

extern uintptr_t kernelEnd;

PageTablePointer *kernelPageTable;
PagePointer *kernelPages;
uint32_t *allocatedPages;
// uint32_t *freePageChunks; // todo: hierarchical structure to find a free page
// more quickly (make one bit an indicator weacher any
// page in the next layer is free) need to find a compromise between size and
// speed

void *getPage();

void identityMapPagesUntil(void *end) {
    void *currentAddress = (void *)0;
    while (currentAddress < end) {
        currentAddress = getPage();
    }
}

void enablePaging() {
    asm volatile("mov %%eax, %%cr3" : : "a"(kernelPageTable));
    asm volatile("mov %cr0, %eax");
    asm volatile("orl $0x80000000, %eax");
    asm volatile("mov %eax, %cr0");
}

void initializePaging() {
    void *currentPosition = (void *)((kernelEnd - 1) & ~0xFFF) + 0x1000;
    kernelPageTable = currentPosition;
    currentPosition += 0x4000;
    kernelPages = currentPosition;
    currentPosition += sizeof(uint32_t) * 1024 * 1024;
    allocatedPages = currentPosition;
    currentPosition +=
        sizeof(uint32_t) * 1024 * 1024 / 32; // 1024 tables with 1024 pages each
    for (int i = 0; i < 1024; i++) {
        kernelPageTable[i].present = 1;
        kernelPageTable[i].writable = 1;
        kernelPageTable[i].pagePointerAddress =
            (uint32_t)(uintptr_t)(kernelPages + 1024 * i) >> 12;
    }
    for (uint32_t i = 0; i < 0x100000; i++) {
        // mark all as inaccessible
        kernelPages[i].present = 0;
    }
    identityMapPagesUntil(currentPosition);
    enablePaging();
}

void allocatePage(uint32_t index32bit, uint8_t bit) {
    allocatedPages[index32bit] |= 1 << bit;
    uint32_t pageId = index32bit + bit;
    kernelPages[pageId].present = 1;
    kernelPages[pageId].writable = 1;
    kernelPages[pageId].pageId = pageId;
}

void *getPage() {
    // todo: add task handling etc.
    for (uint32_t i = 0; i < 0x100000; i += 32) {
        if (allocatedPages[i] == ~0) {
            continue;
        }
        for (uint8_t bit = 0; bit < 32; bit++) {
            if (allocatedPages[i] & (1 << bit)) {
                continue;
            }
            allocatePage(i, bit);
            return (void *)(uintptr_t)((i + bit) << 12);
        }
    }
    return NULL;
}

void freePage(void *pageAddress) {
    uintptr_t pageId = (uintptr_t)pageAddress >> 12;
    allocatedPages[pageId / 32] &= ~(1 << (pageId % 32));
    kernelPages[pageId].present = 0;
}

void markMMIO(void *address) {
    uintptr_t pageId = (uintptr_t)address >> 12;
    allocatedPages[pageId / 32] |= 1 << (pageId % 32);
    kernelPages[pageId].present = 1;
    kernelPages[pageId].writable = 1;
    kernelPages[pageId].isVolatile = 1;
    kernelPages[pageId].pageId = pageId;
}