#include <kernel.h>
#include <init/boot.h>
#include <ex/malloc.h>
#include <mm/mm.h>
#include <ke/idt.h>

#define KPAGESIZE (4 * KB)
#define UPAGESIZE (2 * MB)


// Page directory pointer offset
typedef ulong    pdpe_t;

// Page directory offset
typedef ulong    pde_t;

// Page table entry
typedef ulong    pte_t;

// paging.asm
void MmLoadPML4(void *);
void MmEnableWriteProtect(void);
void MmDisableWriteProtect(void);

enum
{
    MF_PRESENT      = 1 << 0,
    MF_READWRITE    = 1 << 1,
    MF_USERMODE     = 1 << 2,
    MF_WRITETHR     = 1 << 3,
    MF_CACHEDIS     = 1 << 4,
    MF_ACCESSED     = 1 << 5,
    MF_DIRTY        = 1 << 6,
    MF_HUGE         = 1 << 7,
    MF_NX           = 1 << 31
};

#define RAM_MAX 32
#define NB_4K 150
                // * 2 MB

//-----------

volatile pdpe_t MmPML4[512] __attribute__((__aligned__(KPAGESIZE)));

volatile pde_t MmPDP[512] __attribute__((__aligned__(KPAGESIZE)));

volatile pde_t MmPD[512 * RAM_MAX] __attribute__((__aligned__(KPAGESIZE)));;

volatile pte_t MmPT[512 * NB_4K] __attribute__((__aligned__(KPAGESIZE)));;

volatile ulong MmStackGuards[2] = { 0 };

//
// Creates our new page table structure and loads it
//
void MmInitPaging(void)
{
    extern MemoryMap_t memoryMap;
    ulong phRamSize = memoryMap.freeRamSize + memoryMap.nonfreeRamSize;

    memzero((void *)&MmPML4[0], sizeof(MmPML4));
    memzero((void *)&MmPDP[0], sizeof(MmPDP));
    memzero((void *)&MmPD[0], sizeof(MmPD));
    memzero((void *)&MmPT[0], sizeof(MmPT));

    for (volatile ulong i = 0; i < 512 *  NB_4K; i++) {
        // STACK GUARD PAGE
        if ((ulong)(i*KPAGESIZE) == (ulong)BtLoaderInfo.stackEndAddr) {
            MmPT[i] = ((ulong)(i*KPAGESIZE));
            MmStackGuards[0] = i;
            continue;
        }

        // ENOMEM like
        if ((ulong)(i*KPAGESIZE) > (ulong)phRamSize) {
            break;
        }

        // STACK GARD PAGE
        if ((ulong)(i*KPAGESIZE) == (ulong)BtLoaderInfo.kernelEndAddr) {
            MmPT[i] = ((ulong)(i*KPAGESIZE));
            MmStackGuards[1] = i;
            continue;
        }

        MmPT[i] = ((ulong)(i*KPAGESIZE)) | MF_PRESENT | MF_READWRITE;
    }

    for (volatile ulong i = 0; i <  NB_4K; i++) {
        MmPD[i] = (ulong)(&MmPT[i*512])| MF_PRESENT | MF_READWRITE;
    }

    for (volatile ulong i =  NB_4K; i < 512 * RAM_MAX; i++) {
        // ENOMEM like
        if ((ulong)(i* UPAGESIZE) > (ulong)phRamSize) {
            break;
        }

        MmPD[i] = 0;
        MmPD[i] = ((ulong)(i* UPAGESIZE)) | MF_PRESENT | MF_READWRITE | MF_HUGE;
    }

    for (volatile int i = 0; i < RAM_MAX; i++) {
        MmPDP[i] = (ulong)(&MmPD[i*512])| MF_PRESENT | MF_READWRITE;
    }

    MmPML4[0] = (ulong)(&MmPDP[0])| MF_PRESENT | MF_READWRITE;

    MmLoadPML4((void *)MmPML4);
    DebugLog("\tPaging tables initialized at %p, %p\n", &MmPD, &MmPT);
}

//
// Reloads the page tables
//
void MmReloadPaging(void)
{
    extern MemoryMap_t memoryMap;
    ulong phRamSize = memoryMap.freeRamSize + memoryMap.nonfreeRamSize;

    for (volatile ulong i = 0; i < 512 *  NB_4K; i++) {
        // STACK GUARD PAGE
        if ((ulong)(i*KPAGESIZE) == (ulong)BtLoaderInfo.stackEndAddr) {
            MmPT[i] = ((ulong)(i*KPAGESIZE));
            MmStackGuards[0] = i;
            continue;
        }

        // ENOMEM like
        if ((ulong)(i*KPAGESIZE) > (ulong)phRamSize) {
            break;
        }

        // STACK GARD PAGE
        if ((ulong)(i*KPAGESIZE) == (ulong)BtLoaderInfo.kernelEndAddr) {
            MmPT[i] = ((ulong)(i*KPAGESIZE));
            MmStackGuards[1] = i;
            continue;
        }

        MmPT[i] = ((ulong)(i*KPAGESIZE)) | MF_PRESENT | MF_READWRITE;
    }

    for (volatile ulong i = 0; i <  NB_4K; i++) {
        MmPD[i] = (ulong)(&MmPT[i*512])| MF_PRESENT | MF_READWRITE;
    }

    for (volatile ulong i =  NB_4K; i < 512 * RAM_MAX; i++) {
        // ENOMEM like
        if ((ulong)(i* UPAGESIZE) > (ulong)phRamSize) {
            break;
        }

        MmPD[i] = 0;
        MmPD[i] = ((ulong)(i* UPAGESIZE)) | MF_PRESENT | MF_READWRITE | MF_HUGE;
    }

    DebugLog("Paging tables reloaded at %p, %p\n", &MmPD, &MmPT);
}

//
// Page fault handler
//
static void PagingHandler(ISRFrame_t *regs)
{
    KeStartPanic("[ISR 0x%x] Irrecoverable Kernel Page Fault at %p\n\n"
                     "  Error code : 0x%x (%b)\n\n"
                     "  RIP: %#016lx   CS:  %#016lx   RSP: %#016lx\n"
                     "  SS:  %#016lx   RAX: %#016lx   RBX: %#016lx\n"
                     "  RCX: %#016lx   RDX: %#016lx   RSI: %#016lx\n"
                     "  RDI: %#016lx   RBP: %#016lx   R8:  %#016lx\n"
                     "  R9:  %#016lx   R10: %#016lx   R11: %#016lx\n"
                     "  R12: %#016lx   R13: %#016lx   R14: %#016lx\n"
                     "  R15: %#016lx   CR0: %#016lx   CR2: %#016lx\n"
                     "  CR3: %#016lx   CR4: %#016lx   CR8: %#016lx\n"
                     "  RFLAGS: %#022b (%#06x)",
                     regs->intNo,
                     regs->regs[1],
                     regs->ErrorCode,
                     regs->ErrorCode,
                     regs->rip,
                     regs->cs,
                     regs->rsp,
                     regs->ss,
                     regs->regs[5],
                     regs->regs[6],
                     regs->regs[7],
                     regs->regs[8],
                     regs->regs[9],
                     regs->regs[10],
                     regs->regs[11],
                     regs->regs[12],
                     regs->regs[13],
                     regs->regs[14],
                     regs->regs[15],
                     regs->regs[16],
                     regs->regs[17],
                     regs->regs[18],
                     regs->regs[19],
                     regs->regs[0],
                     regs->regs[1],
                     regs->regs[2],
                     regs->regs[3],
                     regs->regs[4],
                     regs->rflags,
                     regs->rflags
                 );
}

void MmActivatePageHandler(void)
{
    KeRegisterISR(PagingHandler, 0xe);
}