#include <kernel.h>
#include <init/boot.h>
#include <ex/malloc.h>
#include <mm/heap.h>
#include <mm/mm.h>
#include <ke/idt.h>
#include <lib/buf.h>
#include <io/vga.h>

#define USERSPACE 0x80000000

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

pml4_t MmPageMapLevel4[512] __attribute__((__aligned__(KPAGESIZE)));
ulong *MmPhysicalPageTable;

extern MemoryMap_t memoryMap;

extern ulong _text;
extern ulong _text_end;
extern ulong _rodata;
extern ulong _rodata_end;
extern ulong _data;
extern ulong _data_end;

ulong MmStackGuards[2] = { 0 };
ulong MmVirtLastAddress = 0;
ulong MmPhysLastKernelAddress = 0;

enum
{
    PRESENT      = 1 << 0,
    READWRITE    = 1 << 1,
    USERMODE     = 1 << 2,
    WRITETHR     = 1 << 3,
    CACHEDIS     = 1 << 4,
    ACCESSED     = 1 << 5,
    DIRTY        = 1 << 6,
    HUGE         = 1 << 7,
    NX           = 1UL << 63
};

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

//
// Creates our new page table structure and loads it
//
void MmInitPaging(void)
{
    pdpe_t *MmPDP = NULL;
    pde_t *MmPD = NULL;
    pte_t *MmPT = NULL;
    ulong index, xedni;
    ulong firstDirectoryAddr = 0;
    ulong lastDirectoryAddr = 0;
    ulong phDirSize = 0;

    // Maximum PHYSICAL address in memory
    ulong phRamSize = memoryMap.freeRamSize + memoryMap.nonfreeRamSize;

    // Difference between the end of kernel and the begin of userspace
    MmPhysLastKernelAddress = (ulong)(_heap_start + _heap_max);
    ulong diffKernUsr = (ulong)USERSPACE - MmPhysLastKernelAddress - KPAGESIZE;

    // Maximum VIRTUAL address in memory
    MmVirtLastAddress = phRamSize + diffKernUsr;

    //DebugLog("\tPaging gap : %u MB (%p)\n\tLast virtual address %p\n", diffKernUsr / MB, diffKernUsr, MmVirtLastAddress);

    memzero((void *)&MmPageMapLevel4[0], sizeof(MmPageMapLevel4));
    phDirSize = ((phRamSize / KPAGESIZE)*sizeof(ulong) + KPAGESIZE) & ( ~(KPAGESIZE - 1));

    MmPhysicalPageTable = (ulong*)malloc(phDirSize);
    //DebugLog("\t\tRam %u MB, pagesize %u KB, size %u MB\n", phRamSize / MB, KPAGESIZE / KB, phDirSize / MB);

    for (ulong curAddrPML4 = 0;
         curAddrPML4 < MmVirtLastAddress;
         curAddrPML4 += ((ulong)KPAGESIZE * 0x8000000)) {
            // Create an entry in PML4 each 512GB
            // 0x8000000 = 512 ^ 3

            MmPDP = (pdpe_t *)malloc(512*sizeof(pde_t));

            if (!firstDirectoryAddr) {
                firstDirectoryAddr = (ulong)MmPDP;
            }

            index = (curAddrPML4 / ((ulong)KPAGESIZE * 0x8000000)) % 512;

            //DebugLog("\t\t\t\tPDP %d : %p\n", index, MmPDP);
            MmPageMapLevel4[index] = (pdpe_t *)((ulong)MmPDP | PRESENT | READWRITE);

            for (ulong curAddrPDP = curAddrPML4;
                 curAddrPDP < (curAddrPML4 + ((ulong)KPAGESIZE * 0x8000000)) &&
                 curAddrPDP < MmVirtLastAddress;
                 curAddrPDP += ((ulong)KPAGESIZE * 0x40000)) {
                    // Create an intry in PDP each 1GB
                    // 0x40000 = 512 ^ 2

                    MmPD = (pde_t *)malloc(512*sizeof(pde_t));

                    index = (curAddrPDP / ((ulong)KPAGESIZE * 0x40000)) % 512;

                    //DebugLog("\t\t\t\tPD %d : %p\n", index, MmPD);
                    MmPDP[index] = (pde_t *)((ulong)MmPD | PRESENT | READWRITE);

                    for (ulong curAddrPD = curAddrPDP;
                         curAddrPD < (curAddrPDP + ((ulong)KPAGESIZE * 0x40000)) &&
                         curAddrPD < MmVirtLastAddress;
                         curAddrPD += ((ulong)KPAGESIZE * 0x200)) {
                            // Create an intry in PD each 2MB
                            // 0x200 = 512

                            MmPT = (pte_t *)malloc(512*sizeof(pte_t));

                            index = (curAddrPD / ((ulong)KPAGESIZE * 0x200)) % 512;

                            //DebugLog("\t\t\t\tPT %d : %p\n", index, MmPT);
                            MmPD[index] = (pte_t *)((ulong)MmPT | PRESENT | READWRITE);

                            for (ulong curAddrPT = curAddrPD;
                                 curAddrPT < (curAddrPD + ((ulong)KPAGESIZE * 0x200)) &&
                                 curAddrPT < MmVirtLastAddress;
                                 curAddrPT += (ulong)KPAGESIZE) {
                                    // Create an entry in PT each page of 4KB

                                    index = (curAddrPT / ((ulong)KPAGESIZE)) % 512;
                                    xedni = (curAddrPT / ((ulong)KPAGESIZE));

                                    //DebugLog("\t\t\t\tPage %d : %p\n", index, curAddrPT);

                                    // STACK GUARD PAGE */
                                    if ((ulong)curAddrPT == (ulong)BtLoaderInfo.stackEndAddr) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;
                                        MmStackGuards[0] = (ulong)curAddrPT;
                                        //DebugLog("\tStack Guard at %p\n", curAddrPT);
                                    }
                                    else if ((ulong)curAddrPT == (ulong)BtLoaderInfo.kernelEndAddr) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;
                                        MmStackGuards[1] = (ulong)curAddrPT;
                                        //DebugLog("\tStack Guard at %p\n", curAddrPT);
                                    }
                                    // SECTION .TEXT PROTECTION
                                    else if ((ulong)curAddrPT >= (ulong)&_text && (ulong)curAddrPT <= (ulong)&_text_end) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;
                                        //DebugLog("\tSection .text at %p\n", curAddrPT);
                                    }
                                    // SECTION .DATA PROTECTION
                                    else if ((ulong)curAddrPT >= (ulong)&_data && (ulong)curAddrPT <= (ulong)&_data_end) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT | WRITETHR | READWRITE | NX;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;
                                        //DebugLog("\tSection .data at %p\n", curAddrPT);
                                    }
                                    // SECTION .RODATA PROTECTION
                                    else if ((ulong)curAddrPT >= (ulong)&_rodata && (ulong)curAddrPT <= (ulong)&_rodata_end) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT | WRITETHR | NX;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;
                                        //DebugLog("\tSection .rodata at %p\n", curAddrPT);
                                    }
                                    // While we're inside the kernel pages
                                    else if ((ulong)curAddrPT <= MmPhysLastKernelAddress) {
                                        MmPT[index] = (ulong)curAddrPT | PRESENT | READWRITE;
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;

                                        if ((ulong)curAddrPT == MmPhysLastKernelAddress) {
                                            //DebugLog("\tLast page of kernel at %p\n", curAddrPT);
                                        }
                                    }
                                    // While we're inside the userspace pages
                                    else if ((ulong)curAddrPT >= USERSPACE) {
                                        MmPT[index] = ((ulong)curAddrPT - diffKernUsr); // Not present for instance, unallocated
                                        xedni = (((ulong)curAddrPT - diffKernUsr) / ((ulong)KPAGESIZE));
                                        MmPhysicalPageTable[xedni] = (ulong)curAddrPT;

                                        if ((ulong)curAddrPT == USERSPACE) {
                                            //DebugLog("\tUserspace at %p:%p\n", curAddrPT, curAddrPT - diffKernUsr);
                                        }
                                    }
                                    else {
                                        MmPT[index] = 0;
                                    }

                                    KeFlushTlbSingle(curAddrPT);
                            }
                    }
            }
    }
    lastDirectoryAddr = (ulong)MmPT;

    MmLoadPML4((void *)MmPageMapLevel4);
    //MmEnableWriteProtect();

    DebugLog("\tPage table size : %u MB\n", (lastDirectoryAddr - firstDirectoryAddr + phDirSize)/MB);
}

//
// Get a page from an address
//
static pte_t *MmGetPageDescriptorFromVirtual(void *virtualAddr)
{
    ulong virtAddrPage = (ulong)virtualAddr & ( ~(KPAGESIZE - 1));

    if (virtAddrPage > MmVirtLastAddress) {
        KeStartPanic("MmSetPage() Out of bound of the address space !");
    }

    pdpe_t *pdp = (pdpe_t*)((ulong)MmPageMapLevel4[(virtAddrPage / ((ulong)KPAGESIZE * 0x8000000)) % 512] & ~(KPAGESIZE - 1));
    //DebugLog("pdp\t: %p\n", pdp);
    pde_t *pd = (pde_t*)( (ulong)pdp[(virtAddrPage / ((ulong)KPAGESIZE * 0x40000)) % 512] & ~(KPAGESIZE - 1));
    //DebugLog("pd\t: %p\n", pd);
    pte_t *pt = (pte_t*)( (ulong)pd[(virtAddrPage / ((ulong)KPAGESIZE * 0x200)) % 512] & ~(KPAGESIZE - 1));
    //DebugLog("pt\t: %p\n", pt);

    pte_t *page = &pt[(virtAddrPage / ((ulong)KPAGESIZE)) % 512];
    //DebugLog("page (with flags): %p\n", page);

    return page;
}


//
// Translates a virtual address to its physical equivalent
//
void *MmTransVirtToPhyAddr(void* virtualAddr)
{
    ulong virtAddrPage = (ulong)virtualAddr & ( ~(KPAGESIZE - 1));
    pte_t *page = MmGetPageDescriptorFromVirtual(virtualAddr);

    if (*page & PRESENT) {
        return NULL;
    }

    return (void*)((*page & ~(KPAGESIZE - 1))+ ((ulong)virtualAddr - (ulong)virtAddrPage));
}

void *MmTransPhyToVirtAddr(void* physicalAddr)
{
    ulong phyAddrPage = (ulong)physicalAddr & ( ~(KPAGESIZE - 1));
    return (void*)( MmPhysicalPageTable[(ulong)physicalAddr
                   / ((ulong)KPAGESIZE)
                   + ((ulong)physicalAddr - phyAddrPage) ] );
}

//
// Add flags to a page
//
void MmSetPage(void* virtualAddr, ulong flags)
{
    pte_t *page = MmGetPageDescriptorFromVirtual(virtualAddr);

    *page |= flags;

    KeFlushTlbSingle(*page);
}

//
// Remove flags of a page
//
void MmUnsetPage(void* virtualAddr, ulong flags)
{
    pte_t *page = MmGetPageDescriptorFromVirtual(virtualAddr);

    *page &= (~flags);

    KeFlushTlbSingle(*page);
}

//
// Map a page in memory
//
void MmMapPage(void* virtualAddr, void* physicalAddr, ulong flags)
{
    pte_t *page = MmGetPageDescriptorFromVirtual(virtualAddr);

    *page = ((ulong)physicalAddr & ~(KPAGESIZE - 1)) | flags;

    KeFlushTlbSingle(*page);
}

//
// Unmap a page in memory
//
void MmUnmapPage(void* virtualAddr)
{
    pte_t *page = MmGetPageDescriptorFromVirtual(virtualAddr);

    *page &= (~PRESENT);

    KeFlushTlbSingle(*page);
}

//
// Find physical unallocated pages
//
void *MmGetPhyPageBlock(size_t size, bool usermode) {
    void *startPhyPage = 0;
    void *endPhyPage = 0;
    size_t curSize = 0;

    // Maximum PHYSICAL address in memory
    ulong phRamSize = memoryMap.freeRamSize + memoryMap.nonfreeRamSize;

    ulong curVirtAddr = 0;

    if (!usermode) {
        startPhyPage = MmTransVirtToPhyAddr((void*)KPAGESIZE + 1);
        endPhyPage = (void*)MmPhysLastKernelAddress;
    } else {
        startPhyPage = MmTransVirtToPhyAddr((void*)USERSPACE);
        endPhyPage = (void*)(phRamSize & ~(KPAGESIZE - 1));
    }

    DebugLog("Start phy at %p\n", startPhyPage);
    DebugLog("End phy at  %p\n", endPhyPage);

    for (ulong curPhyAddr = (ulong)startPhyPage;
         curPhyAddr <= (ulong)endPhyPage;
         curPhyAddr += KPAGESIZE) {
             curVirtAddr = (ulong)MmTransPhyToVirtAddr((void*)curPhyAddr);

             if (curVirtAddr == 0) {
                 DebugLog("CurrentAddr %p\n", curVirtAddr);
                 DebugLog("\t Free !\n");
                 curSize += KPAGESIZE;
             }

             if (curSize >= size)
                 break;
         }

    return NULL;
}

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

//
// Returns the rank of the Stack Guards
//
void *MmGetStackGuards(char rank)
{
    return (void *)MmStackGuards[(int)rank];
}

//
// Page fault handler
//
static void PagingHandler(ISRFrame_t *regs)
{
    ulong StackGuardOne = (ulong)MmGetStackGuards(0);
    ulong StackGuardTwo = (ulong)MmGetStackGuards(1);
    if ((regs->cr2 >= StackGuardOne) && (regs->cr2 <= StackGuardOne + KPAGESIZE) && (regs->rsp <= regs->cr2)) {
        bprintf(BStdOut,
            "\n\n%CPANIC\n[ISR 0x8] Irrecoverable Kernel Stack Underflow\n\n"
                     "  Page Fault Error code : %#x (%b)\n"
                     "  Stack Guard bypassed    : %#x",

                     VGA_COLOR_LIGHT_RED,
                     regs->ErrorCode,
                     regs->ErrorCode,
                     StackGuardOne
                 );
    } else if ((regs->cr2 >= StackGuardTwo) && (regs->cr2 <= StackGuardTwo + KPAGESIZE) && (regs->rsp >= regs->cr2)) {
        bprintf(BStdOut,
            "\n\n%CPANIC\n[ISR 0x8] Irrecoverable Kernel Stack Overflow\n\n"
                     "  Page Fault Error code : %#x (%b)\n"
                     "  Stack Guard bypassed    : %#x",

                     VGA_COLOR_LIGHT_RED,
                     regs->ErrorCode,
                     regs->ErrorCode,
                     StackGuardTwo
                 );
    } else if (regs->cr2 == 0) {
        bprintf(BStdOut,
            "\n\n%CPANIC\n[ISR 0x8] Null vector exception !\n\n"
                     "  Page Fault Error code : %#x (%b)\n",

                     VGA_COLOR_LIGHT_RED,
                     regs->intNo,
                     regs->ErrorCode,
                     regs->ErrorCode
                 );
    } else if (regs->cr2 >= MmVirtLastAddress || regs->cr2 <= 0) {
        bprintf(BStdOut,
            "\n\n%CPANIC\n[ISR 0x8] Out of bound of the address space at %p !\n\n"
                     "  End of the address space : %p\n"
                     "  Page Fault Error code : %#x (%b)\n",

                     VGA_COLOR_LIGHT_RED,
                     regs->cr2,
                     MmVirtLastAddress,
                     regs->ErrorCode,
                     regs->ErrorCode
                 );
    } else {
        //XXX page fault
        bprintf(BStdOut, "\n\n%CPANIC\n[ISR 0x8] Irrecoverable Page Fault at %p\n\n"
                     "  Error code : 0x%x (%b)",

                     VGA_COLOR_LIGHT_RED,
                     regs->cr2,
                     regs->ErrorCode,
                     regs->ErrorCode
                 );
    }

    bprintf(BStdOut, "\n  Description : ");

    if (regs->ErrorCode & PRESENT) {
        bprintf(BStdOut, "Page-protection violation ");
    } else {
        bprintf(BStdOut, "Non present page ");
    }
    if (regs->ErrorCode & READWRITE) {
        bprintf(BStdOut, "during write access ");
    } else {
        bprintf(BStdOut, "during read access ");
    }
    if (regs->ErrorCode & (1 << 3))
        bprintf(BStdOut, "from userspace ");
    if (regs->ErrorCode & (1 << 4))
        bprintf(BStdOut, "after instruction fetching ");

    KeBrkDumpRegisters(regs);

    BStdOut->flusher(BStdOut);

    KeHaltCPU();
}

void MmActivatePageHandler(void)
{
    KeRegisterISR(PagingHandler, 0xe);
    DebugLog("\tPaging activated\n");
}