#include "paging.h"
#include "asm.h"
#include "pfa.h"
#include "bootscreen_output.h"

#define putline printInt(__LINE__); \
                putc('\n')

#ifndef SWAPFILE_NOT_REQUIRED
// Include the filesystem driver
#endif

#include "utils.h"

#define SET(field) field |= 0b1
#define CLEAR(field) field &= 0b0


typedef struct {
    uint8_t present : 1;
    uint8_t read_write : 1;
    uint8_t user_supervisor : 1;
    uint8_t cache_policy : 1;
    uint8_t cache_disable : 1;
    uint8_t accessed : 1;
    uint8_t reserved : 3; // Must be 0b000
    uint8_t avl_3 : 3;
    uint64_t pdp_addr : 40;
    uint16_t avl_11 : 11;
    uint8_t no_execute : 1;
} PML4entry;

typedef struct {
    uint8_t present : 1;
    uint8_t read_write : 1;
    uint8_t user_supervisor : 1;
    uint8_t cache_policy : 1;
    uint8_t cache_disable : 1;
    uint8_t accessed : 1;
    uint8_t mbz : 3; // Must be 0 to also ensure it corresponds to 4kb paging
    uint8_t avl_3 : 3;
    uint64_t pd_addr : 40;
    uint16_t avl_11 : 11;
    uint8_t no_execute : 1;
} PDPentry;

typedef struct {
    uint8_t present : 1;
    uint8_t read_write : 1;
    uint8_t user_supervisor : 1;
    uint8_t cache_policy : 1;
    uint8_t cache_disable : 1;
    uint8_t accessed : 1;
    uint8_t mbz : 3; // Must be 0 to also ensure it corresponds to 4kb paging
    uint8_t avl_3 : 3;
    uint64_t pt_addr : 40;
    uint16_t avl_11 : 11;
    uint8_t no_execute : 1;

} PDentry;

typedef struct {
    uint8_t present : 1;
    uint8_t read_write : 1;
    uint8_t user_supervisor : 1;
    uint8_t cache_policy : 1;
    uint8_t cache_disable : 1;
    uint8_t accessed : 1;
    uint8_t dirty : 1;
    uint8_t pat : 1;
    uint8_t global : 1;
    uint8_t avl_3 : 3;
    uint64_t page_base_addr : 40;
    uint16_t avl_11 : 11;
    uint8_t no_execute : 1;
} PTentry;

#define GET_PGTADDR(x) x & 0b0000000000001111111111111111111111111111111111111111000000000000

/* Initialised from EFI's tables */
static PML4entry* PML4tab;
static PDPentry* PDPtab;
static PDentry* PDtab;
static PTentry* PTtab;

static void* make_canonical(void* addr) {
    size_t naddr = (size_t)addr;
    if (naddr & ((size_t)1 << 47))
        naddr |= 0xFFFF000000000000;
    return (void*)naddr;
}

static void* decanonical(void* addr) {
    return (void*)((size_t)addr & 0x0000FFFFFFFFFFFF);
}

void paging_init() {
    /* Find the page tables EFI has set up. Each entry points to the lower page table */
    PML4tab = (PML4entry*)pfa_alloc();
    PDPtab = (PDPentry*)pfa_alloc();
    PDtab = (PDentry*)pfa_alloc();
    PTtab = (PTentry*)pfa_alloc();

    void* pml4base = read_cr3();
    uint64_t cr0 = read_cr0();
    write_cr0(cr0 & (SIZE_MAX - (1 << 16))); // Disable WP and paging for now
    void* pdpbase;
    void* pdbase;
    void* ptbase;

    uint8_t req_pdp = 0, req_pd = 0, req_pt = 0;

    /* You have to pfa_alloc for the page tables or initialise them, then relocate the entries. */
    /* We can wipe all the bits - we have control of the page tables now. */

    for (uint16_t i = 0; i < 512; i++)
    {
        uint64_t pml4e = *((uint64_t*)pml4base + i);
        if ((pml4e & 0x1) && !req_pdp)
        {
            printLn("Valid entry in the pml4 table found");
            pdpbase = (void*)(GET_PGTADDR(pml4e));
            req_pdp = 1;
        }
        PML4entry pml4_new_e;
        pml4_new_e.present = pml4e & 0b1;
        pml4_new_e.read_write = (pml4e >> 1) & 0b1;
        pml4_new_e.user_supervisor = (pml4e >> 2) & 0b1;
        pml4_new_e.cache_policy = (pml4e >> 3) & 0b1;
        pml4_new_e.cache_disable = (pml4e >> 4) & 0b1;
        pml4_new_e.accessed = (pml4e >> 5) & 0b1;
        pml4_new_e.reserved = 0;
        pml4_new_e.avl_3 = 0;
        pml4_new_e.pdp_addr = (uint64_t)PDPtab >> 12;
        pml4_new_e.avl_11 = 0;
        pml4_new_e.no_execute = (pml4e >> 63) & 0b1;

        PML4tab[i] = pml4_new_e;
        tlbflush(PML4tab);
    }

    printLn("Finished writing in PML4 entries");

    for (int i = 0; i < 512; i++)
    {
        if (PML4tab[i].present) {
            printLn("Found present section");
            printInt(PML4tab[i].pdp_addr);
        }
    }

    for (uint16_t i = 0; i < 512; i++)
    {
        uint64_t pdpe = *((uint64_t*)pdpbase + i);
        if ((pdpe & 0x1) && !req_pd)
        {
            printLn("Valid entry in the pdp table found");
            pdbase = (void*)(GET_PGTADDR(pdpe));
            req_pd = 1;
        }
        PDPentry pdp_new_e;
        pdp_new_e.present = pdpe & 0b1;
        pdp_new_e.read_write = (pdpe >> 1) & 0b1;
        pdp_new_e.user_supervisor = (pdpe >> 2) & 0b1;
        pdp_new_e.cache_policy = (pdpe >> 3) & 0b1;
        pdp_new_e.cache_disable = (pdpe >> 4) & 0b1;
        pdp_new_e.accessed = (pdpe >> 5) & 0b1;
        pdp_new_e.mbz = 0b000;
        pdp_new_e.avl_3 = 0;
        pdp_new_e.pd_addr = (uint64_t)PDtab >> 12;
        pdp_new_e.avl_11 = 0;
        pdp_new_e.no_execute = (pdpe >> 63) & 0b1;

        PDPtab[i] = pdp_new_e;
        tlbflush(PDPtab);
    }

    printLn("Finished writing in PDP entries");

    for (uint16_t i = 0; i < 512; i++)
    {
        uint64_t pde = *((uint64_t*)pdbase + i);
        if ((pde & 0x1) && !req_pt)
        {
            printLn("Valid entry in the pd table found");
            ptbase = (void*)(GET_PGTADDR(pde));
            req_pt = 1;
        }
        PDentry pd_new_e;
        pd_new_e.present = pde & 0b1;
        pd_new_e.read_write = (pde >> 1) & 0b1;
        pd_new_e.user_supervisor = (pde >> 2) & 0b1;
        pd_new_e.cache_policy = (pde>> 3) & 0b1;
        pd_new_e.cache_disable = (pde >> 4) & 0b1;
        pd_new_e.accessed = (pde >> 5) & 0b1;
        pd_new_e.mbz = 0b000;
        pd_new_e.avl_3 = 0;
        pd_new_e.pt_addr = (uint64_t)PTtab >> 12;
        pd_new_e.avl_11 = 0;
        pd_new_e.no_execute = (pde >> 63) & 0b1;

        PDtab[i] = pd_new_e;
        tlbflush(PDtab);
    }

    printLn("Finished writing PD entries");

    for (uint16_t i = 0; i < 512; i++)
    {
        uint64_t pte = *((uint64_t*)ptbase + i);

        PTentry pt_new_e;
        pt_new_e.present = pte & 0b1;
        pt_new_e.read_write = (pte >> 1) & 0b1;
        pt_new_e.user_supervisor = (pte >> 2) & 0b1;
        pt_new_e.cache_policy = (pte >> 3) & 0b1;
        pt_new_e.cache_disable = (pte >> 4) & 0b1;
        pt_new_e.accessed = (pte >> 5) & 0b1;
        pt_new_e.dirty = (pte >> 6) & 0b1;
        pt_new_e.pat = (pte >> 7) & 0b1;
        pt_new_e.global = (pte >> 8) & 0b1;
        pt_new_e.avl_3 = 0b000;

        if (pt_new_e.present)
            pt_new_e.page_base_addr = GET_PGTADDR(pte) >> 12;
        else
            pt_new_e.page_base_addr = 0;

        pt_new_e.avl_11 = 0;
        pt_new_e.no_execute = (pte >> 63) & 0b1;

        PTtab[i] = pt_new_e;
        tlbflush(PTtab);
    }

    printLn("Finished writing in PT entries");

    pfa_free((paddr_t)pml4base);
    pfa_free((paddr_t)pdpbase);
    pfa_free((paddr_t)pdbase);
    pfa_free((paddr_t)ptbase);

    printLn("Freed previous page tables");
    printIntBase((uint64_t)PML4tab, 16);
    write_cr3((void*)PML4tab);
    write_cr0(cr0);
    printLn("PML4 base address has been written. Paging is initialised");
}

void paging_map(void* vaddr, size_t len, size_t attrset)
{
    putline;
    while (len)
    {
        putline;
        uint16_t pml4offset = ((size_t)decanonical(vaddr) >> 39) & 0x1FF;
        uint16_t pdpoffset = ((size_t)decanonical(vaddr) >> 30) & 0x1FF;
        uint16_t pdoffset = ((size_t)decanonical(vaddr) >> 21) & 0x1FF;
        uint16_t ptoffset = ((size_t)decanonical(vaddr) >> 12) & 0x1FF;
        uint16_t ppoffset = (size_t)decanonical(vaddr) & 0x1FF;

        putline;
        PML4entry pml4entry = PML4tab[pml4offset];
        PDPentry pdpentry = PDPtab[pdpoffset];
        PDentry pdentry = PDtab[pdoffset];
        PTentry ptentry = PTtab[ptoffset];

        putline;
        if (!(ptentry.present & 0x1)) {
            /* Page is present, now we have to swap to the pagefile. */

#ifndef SWAPFILE_NOT_REQUIRED
            //Swap to the pagefile
#endif
        }
        else {
            putline;
            //Page is not present, make a new page.
            ptentry.page_base_addr = (pfa_alloc() >> 12);
        }
        putline;
        pdentry.pt_addr = ((uint64_t)PTtab >> 12);
        pdpentry.pd_addr = ((uint64_t)PDtab >> 12);
        pml4entry.pdp_addr = ((uint64_t)PDPtab >> 12);
        ptentry.present, pdentry.present, pdpentry.present, pml4entry.present = 0b1; // Set the present bit

        if (attrset >> 1) //Read/write => read/write set
        {
            SET(ptentry.read_write);
            SET(pdentry.read_write);
            SET(pdpentry.read_write);
            SET(pml4entry.read_write);
        }
        else // read only set
        {
            CLEAR(ptentry.read_write);
            CLEAR(pdentry.read_write);
            CLEAR(pdpentry.read_write);
            CLEAR(pml4entry.read_write);
        }
        if (attrset >> 2) // User/Supervisor => user and supervisor set
        {
            SET(ptentry.user_supervisor);
            SET(pdentry.user_supervisor);
            SET(pdpentry.user_supervisor);
            SET(pml4entry.user_supervisor);
        }
        else // User only set
        {
            CLEAR(ptentry.user_supervisor);
            CLEAR(pdentry.user_supervisor);
            CLEAR(pdpentry.user_supervisor);
            CLEAR(pml4entry.user_supervisor);
        }
        if (attrset >> 3) // Writethrough/writeback => writethrough set
        {
            SET(ptentry.cache_policy);
            SET(pdentry.cache_policy);
            SET(pdpentry.cache_policy);
            SET(pml4entry.cache_policy);
        }
        else // writeback set
        {
            CLEAR(ptentry.cache_policy);
            CLEAR(pdentry.cache_policy);
            CLEAR(pdpentry.cache_policy);
            CLEAR(pml4entry.cache_policy);
        }
        if (attrset >> 4) // Cache/No cache => no cache set
        {
            SET(ptentry.cache_disable);
            SET(pdentry.cache_disable);
            SET(pdpentry.cache_disable);
            SET(pml4entry.cache_disable);
        }
        else // cache set
        {
            CLEAR(ptentry.cache_disable);
            CLEAR(pdentry.cache_disable);
            CLEAR(pdpentry.cache_disable);
            CLEAR(pml4entry.cache_disable);
        }
        if (attrset >> 5) // No execute => no execute set
        {
            SET(ptentry.no_execute);
            SET(pdentry.no_execute);
            SET(pdpentry.no_execute);
            SET(pml4entry.no_execute);
        }
        else // execute set
        {
            CLEAR(ptentry.no_execute);
            CLEAR(pdentry.no_execute);
            CLEAR(pdpentry.no_execute);
            CLEAR(pml4entry.no_execute);
        }
        putline;

        PML4tab[pml4offset] = pml4entry;
        PDPtab[pdpoffset] = pdpentry;
        PDtab[pdoffset] = pdentry;
        PTtab[ptoffset] = ptentry;

        putline;
        tlbflush(PML4tab);
        tlbflush(PDPtab);
        tlbflush(PDtab);
        tlbflush(PTtab);
        putline;
        memory_barrier();

        putline;
        len -= EFI_PAGE_SIZE;
        size_t v_t_addr = (size_t)vaddr;
        v_t_addr += EFI_PAGE_SIZE;
        vaddr = (void*)v_t_addr;
        putline;
    }
}

void set_paging_attributes(void* vaddr, size_t len, size_t attrset)
{
    while (len)
    {
        uint16_t pml4offset = ((size_t)decanonical(vaddr) >> 39) & 0x1FF;
        uint16_t pdpoffset = ((size_t)decanonical(vaddr) >> 30) & 0x1FF;
        uint16_t pdoffset = ((size_t)decanonical(vaddr) >> 21) & 0x1FF;
        uint16_t ptoffset = ((size_t)decanonical(vaddr) >> 12) & 0x1FF;
        uint16_t ppoffset = (size_t)decanonical(vaddr) & 0x1FF;

        PML4entry pml4entry = PML4tab[pml4offset];
        PDPentry pdpentry = PDPtab[pdpoffset];
        PDentry pdentry = PDtab[pdoffset];
        PTentry ptentry = PTtab[ptoffset];

        pdentry.pt_addr = ((uint64_t)PTtab >> 12);
        pdpentry.pd_addr = ((uint64_t)PDtab >> 12);
        pml4entry.pdp_addr = ((uint64_t)PDPtab >> 12);
        ptentry.present, pdentry.present, pdpentry.present, pml4entry.present = 0b1; // Set the present bit

        if (attrset >> 1) //Read/write => read/write set
        {
            SET(ptentry.read_write);
            SET(pdentry.read_write);
            SET(pdpentry.read_write);
            SET(pml4entry.read_write);
        }
        else // read only set
        {
            CLEAR(ptentry.read_write);
            CLEAR(pdentry.read_write);
            CLEAR(pdpentry.read_write);
            CLEAR(pml4entry.read_write);
        }
        if (attrset >> 2) // User/Supervisor => user and supervisor set
        {
            SET(ptentry.user_supervisor);
            SET(pdentry.user_supervisor);
            SET(pdpentry.user_supervisor);
            SET(pml4entry.user_supervisor);
        }
        else // User only set
        {
            CLEAR(ptentry.user_supervisor);
            CLEAR(pdentry.user_supervisor);
            CLEAR(pdpentry.user_supervisor);
            CLEAR(pml4entry.user_supervisor);
        }
        if (attrset >> 3) // Writethrough/writeback => writethrough set
        {
            SET(ptentry.cache_policy);
            SET(pdentry.cache_policy);
            SET(pdpentry.cache_policy);
            SET(pml4entry.cache_policy);
        }
        else // writeback set
        {
            CLEAR(ptentry.cache_policy);
            CLEAR(pdentry.cache_policy);
            CLEAR(pdpentry.cache_policy);
            CLEAR(pml4entry.cache_policy);
        }
        if (attrset >> 4) // Cache/No cache => no cache set
        {
            SET(ptentry.cache_disable);
            SET(pdentry.cache_disable);
            SET(pdpentry.cache_disable);
            SET(pml4entry.cache_disable);
        }
        else // cache set
        {
            CLEAR(ptentry.cache_disable);
            CLEAR(pdentry.cache_disable);
            CLEAR(pdpentry.cache_disable);
            CLEAR(pml4entry.cache_disable);
        }
        if (attrset >> 5) // No execute => no execute set
        {
            SET(ptentry.no_execute);
            SET(pdentry.no_execute);
            SET(pdpentry.no_execute);
            SET(pml4entry.no_execute);
        }
        else // execute set
        {
            CLEAR(ptentry.no_execute);
            CLEAR(pdentry.no_execute);
            CLEAR(pdpentry.no_execute);
            CLEAR(pml4entry.no_execute);
        }

        PML4tab[pml4offset] = pml4entry;
        PDPtab[pdpoffset] = pdpentry;
        PDtab[pdoffset] = pdentry;
        PTtab[ptoffset] = ptentry;

        tlbflush(PML4tab);
        tlbflush(PDPtab);
        tlbflush(PDtab);
        tlbflush(PTtab);
        memory_barrier();

        len -= EFI_PAGE_SIZE;
        size_t v_t_addr = (size_t)vaddr;
        v_t_addr += EFI_PAGE_SIZE;
        vaddr = (void*)v_t_addr;
    }
}