#define MODULE "vmm"

#include "memory/VMM.h"
#include "assert.h"
#include "log/Log.h"
#include "memory/PMM.h"
#include "misc/utils.h"
#include "std/string.h"
#include "utils/Addresses.h"
#include "utils/Registers.h"

static PageTable* kernel_pml4;
static PageTable* current_pml4;
static AddressSpace* user_address_space;

void VMM::switch_back_to_kernel_address_space()
{
    if (current_pml4 != kernel_pml4) { current_pml4 = kernel_pml4; }
}

void VMM::switch_to_user_address_space(AddressSpace& space)
{
    user_address_space = &space;
    current_pml4 = user_address_space->get_pml4();
}

void VMM::switch_to_previous_user_address_space()
{
    current_pml4 = user_address_space->get_pml4();
}

void VMM::enter_syscall_context()
{
    if (current_pml4 != kernel_pml4)
    {
        current_pml4 = kernel_pml4;
        apply_address_space();
        switch_to_previous_user_address_space();
    }
}

void VMM::exit_syscall_context()
{
    if (current_pml4 != user_address_space->get_pml4()) { switch_to_previous_user_address_space(); }
    apply_address_space();
}

void VMM::apply_address_space()
{
    write_cr3(current_pml4);
}

bool VMM::is_using_kernel_address_space()
{
    return current_pml4 == kernel_pml4;
}

void VMM::init()
{
    kernel_pml4 = (PageTable*)read_cr3();
    current_pml4 = kernel_pml4;
}

void VMM::unmap(uint64_t vaddr)
{
    vaddr = round_down_to_nearest_page(vaddr);

    PageDirectoryEntry* pde = find_pde(current_pml4, vaddr);
    if (!pde) return; // Already unmapped

    memset(pde, 0, sizeof(PageDirectoryEntry));
    flush_tlb(vaddr);
}

void VMM::remap(uint64_t vaddr, int flags)
{
    vaddr = round_down_to_nearest_page(vaddr);

    PageDirectoryEntry* pde = find_pde(current_pml4, vaddr);
    if (!pde) return; // Not mapped

    if (flags & User) propagate_user(current_pml4, vaddr);
    else
        pde->user = false;
    if (flags & ReadWrite) propagate_read_write(current_pml4, vaddr);
    else
        pde->read_write = false;
    flush_tlb(vaddr);
}

uint64_t VMM::get_physical(uint64_t vaddr)
{
    PageDirectoryEntry* pde = find_pde(current_pml4, round_down_to_nearest_page(vaddr));
    if (!pde) return UINT64_MAX; // Not mapped

    return pde->get_address() | (vaddr % PAGE_SIZE);
}

uint64_t VMM::get_flags(uint64_t vaddr)
{
    PageDirectoryEntry* pde = find_pde(current_pml4, round_down_to_nearest_page(vaddr));
    if (!pde) return 0; // Not mapped

    uint64_t flags = 0;
    if (pde->user) flags |= User;
    if (pde->read_write) flags |= ReadWrite;
    return flags;
}

void VMM::map(uint64_t vaddr, uint64_t paddr, int flags)
{
    vaddr = round_down_to_nearest_page(vaddr);
    PageDirectoryEntry* pde = find_pde(current_pml4, vaddr);
    bool will_flush_tlb = true;
    if (!pde)
    {
        pde = create_pde_if_not_exists(current_pml4, vaddr);
        will_flush_tlb = false;
    }
    else if (pde->larger_pages)
    {
        unmap(vaddr);
        pde = create_pde_if_not_exists(current_pml4, vaddr);
    }

    pde->set_address(round_down_to_nearest_page(paddr));
    if (flags & User) propagate_user(current_pml4, vaddr);
    else
        pde->user = false;
    if (flags & ReadWrite) propagate_read_write(current_pml4, vaddr);
    else
        pde->read_write = false;
    if (will_flush_tlb) flush_tlb(vaddr);
}

PageDirectoryEntry* VMM::find_pde(PageTable* root, uint64_t vaddr)
{
    uint64_t page_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    uint64_t indexes[3];

    decompose_vaddr(vaddr, page_index, indexes[2], indexes[1], indexes[0]);

    for (int i = 0; i < 3;
         i++) // Walk through the page map level 4, page directory pointer, and page directory to find the page table.
    {
        pde = &pt->entries[indexes[i]];
        if (!pde->present) return nullptr;
        else if (pde->larger_pages)
            return pde;
        else { pt = (PageTable*)pde->get_address(); }
    }

    pde = &pt->entries[page_index]; // PT
    if (!pde->present) return nullptr;
    return pde;
}

PageDirectoryEntry* VMM::create_pde_if_not_exists(PageTable* root, uint64_t vaddr)
{
    uint64_t page_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    uint64_t indexes[3];

    decompose_vaddr(vaddr, page_index, indexes[2], indexes[1], indexes[0]);

    auto pde_create_if_not_present = [&]() {
        pt = (PageTable*)PMM::request_page();
        ASSERT(!(PMM_DID_FAIL(pt)));
        memset(pt, 0, PAGE_SIZE);
        pde->set_address((uint64_t)pt);
        pde->present = true;
    };

    for (int i = 0; i < 3; i++)
    {
        pde = &pt->entries[indexes[i]];
        if (!pde->present) { pde_create_if_not_present(); }
        else if (pde->larger_pages)
            return pde;
        else { pt = (PageTable*)pde->get_address(); }
    }

    pde = &pt->entries[page_index];
    if (!pde->present) { pde->present = true; }
    return pde;
}

void VMM::propagate_read_write(PageTable* root, uint64_t vaddr)
{
    uint64_t page_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    uint64_t indexes[3];

    decompose_vaddr(vaddr, page_index, indexes[2], indexes[1], indexes[0]);

    for (int i = 0; i < 3; i++)
    {
        pde = &pt->entries[indexes[i]];
        if (!pde->present) return;
        else
        {
            pde->read_write = true;
            if (pde->larger_pages) return;
            pt = (PageTable*)pde->get_address();
        }
    }

    pde = &pt->entries[page_index];
    if (!pde->present) return;
    else
        pde->read_write = true;
}

void VMM::propagate_user(PageTable* root, uint64_t vaddr)
{
    uint64_t page_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    uint64_t indexes[3];

    decompose_vaddr(vaddr, page_index, indexes[2], indexes[1], indexes[0]);

    for (int i = 0; i < 3; i++)
    {
        pde = &pt->entries[indexes[i]];
        if (!pde->present) return;
        else
        {
            pde->user = true;
            if (pde->larger_pages) return;
            pt = (PageTable*)pde->get_address();
        }
    }

    pde = &pt->entries[page_index];
    if (!pde->present) return;
    else
        pde->user = true;
}

void VMM::flush_tlb(uint64_t addr)
{
    asm volatile("invlpg (%0)" : : "r"(addr) : "memory");
}

void VMM::decompose_vaddr(uint64_t vaddr, uint64_t& page_index, uint64_t& pt_index, uint64_t& pd_index,
                          uint64_t& pdp_index)
{
    vaddr >>= 12;
    page_index = vaddr & 0x1ff;
    vaddr >>= 9;
    pt_index = vaddr & 0x1ff;
    vaddr >>= 9;
    pd_index = vaddr & 0x1ff;
    vaddr >>= 9;
    pdp_index = vaddr & 0x1ff;
}

uint64_t VMM::recompose_vaddr(uint64_t page_index, uint64_t pt_index, uint64_t pd_index, uint64_t pdp_index)
{
    return pdp_index << 39 | pd_index << 30 | pt_index << 21 | page_index << 12;
}

void VMM::install_kernel_page_directory_into_address_space(AddressSpace& space)
{
    PageTable* space_pml4 = space.get_pml4();
    PageTable* kernel_last_pdp = (PageTable*)kernel_pml4->entries[511].get_address();
    PageTable* kernel_last_pd = (PageTable*)kernel_last_pdp->entries[511].get_address();

    PageTable* space_last_pdp = (PageTable*)PMM::request_page();

    PageDirectoryEntry& space_last_pdp_pde = space_pml4->entries[511];

    space_last_pdp_pde.present = true;
    space_last_pdp_pde.read_write = true;
    space_last_pdp_pde.set_address((uint64_t)space_last_pdp);

    PageDirectoryEntry& space_last_pd_pde = space_last_pdp->entries[511];

    space_last_pd_pde.present = true;
    space_last_pd_pde.read_write = true;
    space_last_pd_pde.set_address((uint64_t)kernel_last_pd);
}