#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();
    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);
}

AddressSpace AddressSpace::clone() // FIXME: Add out-of-memory checks to this function.
{
    AddressSpace result;
    result.m_pml4 = (PageTable*)PMM::request_page();
    if (!result.m_pml4) return result;
    memcpy(result.m_pml4, m_pml4, PAGE_SIZE);
    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 = PMM::request_page();
            memcpy(cloned, (void*)pdp_pde.get_address(), PAGE_SIZE);
            cloned_pdp_pde.set_address((uint64_t)cloned);
            continue;
        }
        PageTable* pdp = (PageTable*)pdp_pde.get_address();
        PageTable* cloned_pdp = (PageTable*)PMM::request_page();
        memcpy(cloned_pdp, pdp, PAGE_SIZE);
        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 = PMM::request_page();
                memcpy(cloned, (void*)pd_pde.get_address(), PAGE_SIZE);
                cloned_pd_pde.set_address((uint64_t)cloned);
                continue;
            }
            PageTable* pd = (PageTable*)pd_pde.get_address();
            PageTable* cloned_pd = (PageTable*)PMM::request_page();
            memcpy(cloned_pd, pd, PAGE_SIZE);
            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 = PMM::request_page();
                    memcpy(cloned, (void*)pt_pde.get_address(), PAGE_SIZE);
                    cloned_pt_pde.set_address((uint64_t)cloned);
                    continue;
                }
                PageTable* pt = (PageTable*)pt_pde.get_address();
                PageTable* cloned_pt = (PageTable*)PMM::request_page();
                memcpy(cloned_pt, pt, PAGE_SIZE);
                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 = PMM::request_page();
                    memcpy(cloned, (void*)pde.get_address(), PAGE_SIZE);
                    cloned_pde.set_address((uint64_t)cloned);
                    continue;
                }
            }
        }
    }
    return result;
}