#define MODULE "vmm"

#include "memory/AddressSpace.h"
#include "log/Log.h"
#include "memory/PMM.h"
#include "memory/VMM.h"
#include "std/stdlib.h"
#include "std/string.h"
#include "utils/move.h"

AddressSpace AddressSpace::create()
{
    AddressSpace result;
    result.m_pml4 = (PageTable*)PMM::request_page();
    memset(result.m_pml4, 0, PAGE_SIZE);
    VMM::install_kernel_page_directory_into_address_space(result);
    return move(result);
}

void AddressSpace::destroy()
{
    uint64_t pages_freed = 0;
    for (int i = 0; i < 512; i++)
    {
        PageDirectoryEntry& pdp_pde = m_pml4->entries[i];
        if (!pdp_pde.present) continue;
        if (pdp_pde.larger_pages)
        {
            pages_freed++;
            PMM::free_page((void*)pdp_pde.get_address());
            continue;
        }
        PageTable* pdp = (PageTable*)pdp_pde.get_address();
        for (int j = 0; j < 511; j++) // skip the last page directory, it's the kernel one
        {
            PageDirectoryEntry& pd_pde = pdp->entries[j];
            if (!pd_pde.present) continue;
            if (pd_pde.larger_pages)
            {
                pages_freed++;
                PMM::free_page((void*)pd_pde.get_address());
                continue;
            }
            PageTable* pd = (PageTable*)pd_pde.get_address();
            for (int k = 0; k < 512; k++)
            {
                PageDirectoryEntry& pt_pde = pd->entries[k];
                if (!pt_pde.present) continue;
                if (pt_pde.larger_pages)
                {
                    pages_freed++;
                    PMM::free_page((void*)pt_pde.get_address());
                    continue;
                }
                PageTable* pt = (PageTable*)pt_pde.get_address();
                for (int l = 0; l < 512; l++)
                {
                    PageDirectoryEntry& pde = pt->entries[l];
                    if (!pde.present) continue;
                    pages_freed++;
                    PMM::free_page((void*)pde.get_address());
                }
                pages_freed++;
                PMM::free_page(pt);
            }
            pages_freed++;
            PMM::free_page(pd);
        }
        pages_freed++;
        PMM::free_page(pdp);
    }
    pages_freed++;
    PMM::free_page(m_pml4);

    kdbgln("Reclaimed %ld pages from address space!", pages_freed);
}

void AddressSpace::clear()
{
    uint64_t pages_freed = 0;
    for (int i = 0; i < 512; i++)
    {
        PageDirectoryEntry& pdp_pde = m_pml4->entries[i];
        if (!pdp_pde.present) continue;
        if (pdp_pde.larger_pages)
        {
            pages_freed++;
            PMM::free_page((void*)pdp_pde.get_address());
            continue;
        }
        PageTable* pdp = (PageTable*)pdp_pde.get_address();
        for (int j = 0; j < 511; j++) // skip the last page directory, it's the kernel one
        {
            PageDirectoryEntry& pd_pde = pdp->entries[j];
            if (!pd_pde.present) continue;
            if (pd_pde.larger_pages)
            {
                pages_freed++;
                PMM::free_page((void*)pd_pde.get_address());
                continue;
            }
            PageTable* pd = (PageTable*)pd_pde.get_address();
            for (int k = 0; k < 512; k++)
            {
                PageDirectoryEntry& pt_pde = pd->entries[k];
                if (!pt_pde.present) continue;
                if (pt_pde.larger_pages)
                {
                    pages_freed++;
                    PMM::free_page((void*)pt_pde.get_address());
                    continue;
                }
                PageTable* pt = (PageTable*)pt_pde.get_address();
                for (int l = 0; l < 512; l++)
                {
                    PageDirectoryEntry& pde = pt->entries[l];
                    if (!pde.present) continue;
                    pages_freed++;
                    PMM::free_page((void*)pde.get_address());
                }
                pages_freed++;
                PMM::free_page(pt);
            }
            pages_freed++;
            PMM::free_page(pd);
        }
        pages_freed++;
        PMM::free_page(pdp);
    }
    memset(m_pml4, 0, PAGE_SIZE);

    VMM::install_kernel_page_directory_into_address_space(*this);

    kdbgln("Reclaimed %ld pages from address space!", pages_freed);
}

static PageTable* try_clone_page_table(PageTable* source)
{
    PageTable* dst = (PageTable*)PMM::request_page();
    if (PMM_DID_FAIL(dst)) { return 0; }
    memcpy(dst, source, sizeof(PageTable));
    return dst;
}

AddressSpace AddressSpace::clone() // FIXME: Add out-of-memory checks to this function.
{
    AddressSpace result;
    result.m_pml4 = try_clone_page_table(m_pml4);
    if (!result.m_pml4) return result;
    for (int i = 0; i < 512; i++)
    {
        PageDirectoryEntry& pdp_pde = m_pml4->entries[i];
        PageDirectoryEntry& cloned_pdp_pde = result.m_pml4->entries[i];
        if (!pdp_pde.present) continue;
        if (pdp_pde.larger_pages)
        {
            void* cloned = try_clone_page_table((PageTable*)pdp_pde.get_address());
            if (!cloned)
            {
                cloned_pdp_pde.present = false;
                continue;
            }
            cloned_pdp_pde.set_address((uint64_t)cloned);
            continue;
        }
        PageTable* pdp = (PageTable*)pdp_pde.get_address();
        PageTable* cloned_pdp = try_clone_page_table(pdp);
        if (!cloned_pdp)
        {
            cloned_pdp_pde.present = false;
            continue;
        }
        cloned_pdp_pde.set_address((uint64_t)cloned_pdp);
        for (int j = 0; j < 511; j++) // skip the last page directory, it's the kernel one
        {
            PageDirectoryEntry& pd_pde = pdp->entries[j];
            PageDirectoryEntry& cloned_pd_pde = cloned_pdp->entries[j];
            if (!pd_pde.present) continue;
            if (pd_pde.larger_pages)
            {
                void* cloned = try_clone_page_table((PageTable*)pd_pde.get_address());
                if (!cloned)
                {
                    cloned_pd_pde.present = false;
                    continue;
                }
                cloned_pd_pde.set_address((uint64_t)cloned);
                continue;
            }
            PageTable* pd = (PageTable*)pd_pde.get_address();
            PageTable* cloned_pd = try_clone_page_table(pd);
            if (!cloned_pd)
            {
                cloned_pd_pde.present = false;
                continue;
            }
            cloned_pd_pde.set_address((uint64_t)cloned_pd);
            for (int k = 0; k < 512; k++)
            {
                PageDirectoryEntry& pt_pde = pd->entries[k];
                PageDirectoryEntry& cloned_pt_pde = cloned_pd->entries[k];
                if (!pt_pde.present) continue;
                if (pt_pde.larger_pages)
                {
                    void* cloned = try_clone_page_table((PageTable*)pt_pde.get_address());
                    if (!cloned)
                    {
                        cloned_pt_pde.present = false;
                        continue;
                    }
                    cloned_pt_pde.set_address((uint64_t)cloned);
                    continue;
                }
                PageTable* pt = (PageTable*)pt_pde.get_address();
                PageTable* cloned_pt = try_clone_page_table(pt);
                if (!cloned_pt)
                {
                    cloned_pt_pde.present = false;
                    continue;
                }
                cloned_pt_pde.set_address((uint64_t)cloned_pt);
                for (int l = 0; l < 512; l++)
                {
                    PageDirectoryEntry& pde = pt->entries[l];
                    PageDirectoryEntry& cloned_pde = cloned_pt->entries[l];
                    if (!pde.present) continue;
                    void* cloned = try_clone_page_table((PageTable*)pde.get_address());
                    if (!cloned)
                    {
                        cloned_pde.present = false;
                        continue;
                    }
                    cloned_pde.set_address((uint64_t)cloned);
                    continue;
                }
            }
        }
    }
    return result;
}