#include "vmm.h"
#include "pmm.h"

uint32_t page_directory[1024] __attribute__((aligned(4096)));
uint32_t first_page_table[1024] __attribute__((aligned(4096)));

static uint32_t temp_table_phys = 0;

extern void vmm_enable_paging(uint32_t* page_directory);

void vmm_init(void) {
    for (int i = 0; i < 1024; i++) {
        page_directory[i] = 0x00000002;
    }

    for (uint32_t i = 0; i < 1024; i++) {
        first_page_table[i] = (i * 0x1000) | 0x3;
    }

    page_directory[0] = ((uint32_t)first_page_table) | 0x3;

    temp_table_phys = (uint32_t)pmm_alloc_page();
    page_directory[1023] = (temp_table_phys & 0xFFFFF000) | 0x3;

    uint32_t* temp_table = (uint32_t*)temp_table_phys;
    for (int i = 0; i < 1024; i++) temp_table[i] = 0;

    uint32_t bitmap_phys = 0x1000000;
    uint32_t bitmap_pd_idx = bitmap_phys >> 22;
    uint32_t bitmap_pt_idx = (bitmap_phys >> 12) & 0x3FF;
    uint32_t bitmap_pt_phys = (uint32_t)pmm_alloc_page();
    page_directory[bitmap_pd_idx] = (bitmap_pt_phys & 0xFFFFF000) | 0x3;
    uint32_t* bitmap_pt = (uint32_t*)bitmap_pt_phys;
    for (int i = 0; i < 1024; i++) bitmap_pt[i] = 0;
    bitmap_pt[bitmap_pt_idx] = (bitmap_phys & 0xFFFFF000) | 0x3;

    vmm_enable_paging(page_directory);
}

void vmm_map_page(uint32_t phys, uint32_t virt, uint32_t flags) {
    uint32_t pd_idx = virt >> 22;
    uint32_t pt_idx = (virt >> 12) & 0x3FF;
    uint32_t temp_pt_idx = (VMM_TEMP_MAP >> 12) & 0x3FF;
    uint32_t* tmp = (uint32_t*)temp_table_phys;

    if (!(page_directory[pd_idx] & 1)) {
        uint32_t phys_table = (uint32_t)pmm_alloc_page();
        if (!phys_table) {
            return;
        }
        page_directory[pd_idx] = (phys_table & 0xFFFFF000) | 3;

        tmp[temp_pt_idx] = (phys_table & 0xFFFFF000) | 3;
        __asm__ __volatile__("invlpg (%0)" : : "r"(VMM_TEMP_MAP) : "memory");

        uint32_t* new_table = (uint32_t*)VMM_TEMP_MAP;
        for (int i = 0; i < 1024; i++) new_table[i] = 0;

        new_table[pt_idx] = (phys & 0xFFFFF000) | flags;

        tmp[temp_pt_idx] = 0;
        __asm__ __volatile__("invlpg (%0)" : : "r"(VMM_TEMP_MAP) : "memory");
        __asm__ __volatile__("invlpg (%0)" : : "r"(virt) : "memory");
        return;
    }

    uint32_t table_phys = page_directory[pd_idx] & 0xFFFFF000;
    tmp[temp_pt_idx] = table_phys | 3;
    __asm__ __volatile__("invlpg (%0)" : : "r"(VMM_TEMP_MAP) : "memory");

    uint32_t* table = (uint32_t*)VMM_TEMP_MAP;
    table[pt_idx] = (phys & 0xFFFFF000) | flags;

    tmp[temp_pt_idx] = 0;
    __asm__ __volatile__("invlpg (%0)" : : "r"(VMM_TEMP_MAP) : "memory");
    __asm__ __volatile__("invlpg (%0)" : : "r"(virt) : "memory");
}
