#define MODULE "vmm"

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

AddressSpace AddressSpace::create()
{
    AddressSpace result;
    result.m_pml4 = (PageTable*)PMM::request_page();
    result.m_refs = (int*)kmalloc(sizeof(int));
    *result.m_refs = 1;
    VMM::install_kernel_page_directory_into_address_space(result);
    return move(result);
}

void AddressSpace::destroy()
{
    if (is_cloned())
    {
        kdbgln("Will not destroy a cloned address space, I don't own it");
        (*m_refs)--;
        return;
    }
    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);

    kfree(m_refs);

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

void AddressSpace::detach()
{
    if (!is_cloned()) return;
    (*m_refs)--;
    m_refs = (int*)kmalloc(sizeof(int));
    *m_refs = 1;
    m_pml4 = (PageTable*)PMM::request_page();
    VMM::install_kernel_page_directory_into_address_space(*this);
}

AddressSpace AddressSpace::clone()
{
    AddressSpace result;
    result.m_pml4 = m_pml4;
    result.m_refs = m_refs;
    *m_refs = *m_refs + 1;
    return result;
}