#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"

// FIXME: There is a lot of duplicate code in this file. This should probably be refactored.

static PageTable* PML4;

void VMM::init()
{
    asm volatile("mov %%cr3, %0" : "=r"(PML4));
}

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

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

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

uint64_t VMM::get_physical(uint64_t vaddr)
{
    PageDirectoryEntry* pde = find_pde(PML4, Utilities::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(PML4, Utilities::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 = Utilities::round_down_to_nearest_page(vaddr);
    PageDirectoryEntry* pde = find_pde(PML4, vaddr);
    bool will_flush_tlb = true;
    if (!pde)
    {
        pde = create_pde_if_not_exists(PML4, vaddr);
        will_flush_tlb = false;
    }
    else if (pde->larger_pages)
    {
        unmap(vaddr);
        pde = create_pde_if_not_exists(PML4, vaddr);
        will_flush_tlb = false;
    }

    pde->set_address(Utilities::round_down_to_nearest_page(paddr));
    if (flags & User) propagate_user(PML4, vaddr);
    if (flags & ReadWrite) propagate_read_write(PML4, vaddr);
    if (will_flush_tlb) flush_tlb(vaddr);
}

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

    decompose_vaddr(vaddr, page_index, pt_index, pd_index, pdp_index);

    uint64_t indexes[3] = {pdp_index, pd_index, pt_index};

    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, pt_index, pd_index, pdp_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    decompose_vaddr(vaddr, page_index, pt_index, pd_index, pdp_index);

    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;
    };

    uint64_t indexes[3] = {pdp_index, pd_index, pt_index};

    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, pt_index, pd_index, pdp_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    decompose_vaddr(vaddr, page_index, pt_index, pd_index, pdp_index);

    uint64_t indexes[3] = {pdp_index, pd_index, pt_index};

    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, pt_index, pd_index, pdp_index;
    PageDirectoryEntry* pde;
    PageTable* pt = root;

    decompose_vaddr(vaddr, page_index, pt_index, pd_index, pdp_index);

    uint64_t indexes[3] = {pdp_index, pd_index, pt_index};

    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;
}