#include "paging.h"
#include "arch/x86_64/x86_64.h"
#include "hal/memory.h"
#include "jove.h"
#include "proc/init.h"
#include "config.h"
#include "hal/obj.h"
#include "lib/string.h"
#include "lib/printf.h"
#include "panic/panic.h"
#include "idt.h"
#include "api/syscall.h"

PAGEALIGN page_t _kernel_pml4[512];
PAGEALIGN page_t _kernel_pml3[512];
PAGEALIGN page_t _kernel_pml2[512];
PAGEALIGN page_t _kernel_pml1[512];

#define IDENT_GIANTPAGES ((PHYS_HUGEPAGES_MAX + 511) / 512)
#define IDENT_MASSIVEPAGES ((IDENT_GIANTPAGES + 511) / 512)

PAGEALIGN static page_t ident_pml3[IDENT_MASSIVEPAGES][512];
PAGEALIGN static page_t ident_pml2[IDENT_GIANTPAGES][512];

virtaddr_t __syscall_stack_guard;
virtaddr_t __syscall_stack;
virtaddr_t __interrupt_stack_guard;
virtaddr_t __interrupt_stack;

virtaddr_t *__syscall_stack_mapping;
virtaddr_t *__interrupt_stack_mapping;

void*
phys_tovirt(uintptr_t phys)
{
    return (void*)(phys | IDENT_VIRT_BASE);
}

physaddr_t
phys_fromvirt(virtaddr_t virt)
{
    if(virt < IDENT_VIRT_BASE) return 0;
    if(virt > KERNEL_VIRT_BASE) return KPTR_TOPHYS(virt);
    return virt & ~IDENT_VIRT_BASE;
}

uintptr_t
paging_get_pml4(void)
{
    uintptr_t ret = 0;
    __asm__ volatile("movq %%cr3, %0": "=r"(ret));
    return ret;
}

void
paging_pml4_switch(uintptr_t pml4)
{
    __asm__ volatile("movq %0, %%cr3":: "r"(pml4): "memory");
}

void pagedir_invlpg(uintptr_t vptr)
{
    __asm__ volatile("invlpg (%0)":: "r"(vptr));
}

static struct int_stackframe*
s_int_handle_pf_early(struct int_stackframe *frame)
{
    extern int _kernel_panic_unlock;
    _kernel_panic_unlock = 1;

    union {
        struct {
            uint8_t p : 1;
            uint8_t wr : 1;
            uint8_t us : 1;
            uint8_t rsvd : 1;
            uint8_t id : 1;
            uint8_t pk : 1;
        };
        uint64_t value;
    } pf_error = { .value = frame->int_err };
    
    kprintf("Kernel has encountered a page fault"
            "exception early in initializaiton\n");
    kprintf("%s %s during %s\n", 
            pf_error.wr ? "Wrote to" : "Read from",
            pf_error.p ? "Present" : "Non-present",
            pf_error.id ? "Instruction fetch" : "Data access");
    
    uintptr_t cr2 = 0;
    __asm__ volatile("movq %%cr2, %0": "=r"(cr2));
    kpanic_istate(frame, "Faulting address: %p\n", cr2);
}

static page_t*
s_pagedir_get_pml1(page_t *pml4, uintptr_t vptr)
{
    errno = 0;
    size_t pml4i = PML4I(vptr);
    if(!pml4[pml4i].p) {
        errno = ENOMAP;
        return 0;
    }

    #define PML_TBL(tbl, i) (page_t*)(((uintptr_t)tbl[i].pfn << 12) | IDENT_VIRT_BASE)

    page_t *pml3 = PML_TBL(pml4, pml4i);
    size_t pml3i = PML3I(vptr);
    if(!pml3[pml3i].p) {
        errno = ENOMAP;
        return 0;
    }

    page_t *pml2 = PML_TBL(pml3, pml3i); 
    size_t pml2i = PML2I(vptr);
    if(!pml2[pml2i].p) {
        errno = ENOMAP;
        return 0;
    }

    return PML_TBL(pml2, pml2i);
}

int
pagedir_new(jove_owned_obj_t *untyped_obj)
{
    jove_untyped_t *untyped = obj_deref(untyped_obj);
    if((uintptr_t)untyped & 0xFFF) return -EALIGN;
    if(untyped->limit != 0x1000) return -EBADSIZE;

    obj_settype(untyped_obj, JOVE_OBJ_PAGEDIR);
    memcpy(untyped, _kernel_pml4, 0x1000);
    return 0;
}

uintptr_t *pagedir_get_mapping(jove_owned_obj_t *pagedir, uintptr_t vptr)
{
    page_t *pml4 = obj_deref(pagedir);
    page_t *pml1 = s_pagedir_get_pml1(pml4, vptr);
    if(!pml1) return 0;

    size_t pml1i = PML1I(vptr);
    page_t *mapping = &pml1[pml1i];
    return (uintptr_t*)mapping;
}

size_t
pagedir_mapping_get_pfn(virtaddr_t *mapping)
{
    return ((page_t*)mapping)->pfn;
}

virtaddr_t
pagedir_ident_from_mapping(uintptr_t *mapping)
{
    return (((page_t*)mapping)->pfn << 12) | IDENT_VIRT_BASE;
}

inline bool pagedir_mapping_is_present(virtaddr_t *mapping)
{
    return ((page_t*)mapping)->p;
}

inline void pagedir_mapping_set_pfn(virtaddr_t *mapping, size_t pfn)
{
    ((page_t*)mapping)->pfn = pfn;
}

inline void pagedir_mapping_set_present(virtaddr_t *mapping, bool present)
{
    ((page_t*)mapping)->p = present;
}

inline void pagedir_mapping_set_writeable(virtaddr_t *mapping, bool rw)
{
    ((page_t*)mapping)->rw = rw;
}

inline void pagedir_mapping_set_user(virtaddr_t *mapping, bool user)
{
    ((page_t*)mapping)->u = user;
}

inline void pagedir_mapping_set_type(virtaddr_t *mapping, int type)
{
    ((page_t*)mapping)->mapping_type = type;
}

int pagedir_map_frame(page_t *pml4, 
        uintptr_t pptr, uintptr_t vptr,
        uintmax_t flg)
{
    page_t *pml1 = s_pagedir_get_pml1(pml4, vptr);
    if(!pml1)
        return errno;

    size_t pml1i = PML1I(vptr);
    if(pml1[pml1i].p) return EEXIST;

    pml1[pml1i].value = pptr | PAGE_PRESENT | flg;
    pagedir_invlpg(vptr);
    return 0;
}

int pagedir_map_frame_type(page_t *pml4,
        uintptr_t pptr, uintptr_t vptr,
        struct syscall_pagedir_flags flg, int type)
{
    uintmax_t real_flg = PAGE_USER | (type << PAGE_SHL_TYPE);
    if(flg.writeable) {
        real_flg |= PAGE_WRITEABLE;
    }
    if(!flg.executable) {
        real_flg |= PAGE_XD;
    }
    return pagedir_map_frame(pml4, pptr, vptr, real_flg);
}

int pagedir_map_untyped_frame(
        jove_owned_obj_t *tgt, jove_perm_obj_t *untyped, 
        uintptr_t vptr, struct syscall_pagedir_flags flg)
{
    page_t *pml4 = obj_deref(tgt);

    if(flg.writeable && (untyped->permissions & JOVE_PERM_WRITE) == 0) {
        kprintf("%s: Untyped object flagged for writeable mapping missing write permissions!\n", __func__);
        return -EPERM;
    }

    uintptr_t frame = obj_getvalue(untyped->obj);
    return pagedir_map_frame_type(pml4, frame, vptr, flg, JOVE_MAPPING_TYPE_UNTYPED);
}

int pagedir_map_table(page_t *pml4,
        uintptr_t pptr, uintptr_t vptr)
{
    page_t *target_table = NULL;

    void *frame_ident = (void*)(pptr | IDENT_VIRT_BASE);

    size_t pml4i = PML4I(vptr);
    if(!pml4[pml4i].p) {
        target_table = &pml4[pml4i];
        goto map_table;
    }

    page_t *pml3 = (page_t*)(((uintptr_t)pml4[pml4i].pfn << 12)
        | IDENT_VIRT_BASE);
    size_t pml3i = PML3I(vptr);
    if(!pml3[pml3i].p) {
        target_table = &pml3[pml3i];
        goto map_table;
    }

    page_t *pml2 = (page_t*)(((uintptr_t)pml3[pml3i].pfn << 12)
        | IDENT_VIRT_BASE);
    size_t pml2i = PML2I(vptr);
    if(!pml2[pml2i].p) {
        target_table = &pml2[pml2i];
        goto map_table;
    }
    return EHASMAP;

map_table:
    memset(frame_ident, 0, 0x1000);
    target_table->value = pptr |
        PAGE_PRESENT |
        (pptr > IDENT_VIRT_BASE ? 0 : PAGE_USER) |
        PAGE_WRITEABLE;
    pagedir_invlpg(vptr);
    return 0;
}

int pagedir_map_untyped_table(
        jove_owned_obj_t *tgt, jove_perm_obj_t *untyped, uintptr_t vptr)
{
    page_t *pml4 = obj_deref(tgt);

    if((untyped->permissions & JOVE_PERM_WRITE) == 0) {
        kprintf("%s: Untyped object flagged for table mapping missing write permissions!\n", __func__);
        return -EPERM;
    }
    uintptr_t frame = obj_getvalue(untyped->obj);

    return pagedir_map_table(pml4, frame, vptr);
}

int pagedir_map_shmem_frame(
        jove_owned_obj_t *tgt, jove_perm_obj_t *shmem, uint8_t page_offset,
        uintptr_t vptr, struct syscall_pagedir_flags flg)
{
    page_t *pml4 = obj_deref(tgt);

    if(flg.writeable && (shmem->permissions & JOVE_PERM_WRITE) != 0) {
        kprintf("%s: Shared memory object flagged for writeable mapping missing write permissions!\n", __func__);
        return -EPERM;
    }
    
    uintmax_t value = obj_getvalue(shmem->obj);
    uintptr_t frame = value & ~0xFF;
    frame += (uintmax_t)page_offset << 12;

    return pagedir_map_frame_type(pml4, frame, vptr, flg, JOVE_MAPPING_TYPE_SHMEM);
}

int pagedir_unmap_untyped_frame(
        jove_owned_obj_t *tgt, jove_owned_obj_t *dst, uintptr_t vptr)
{
    page_t *pml4 = obj_deref(tgt);

    size_t pml4i = PML4I(vptr);
    if(!pml4[pml4i].p) return -ENOMAP;

    page_t *pml1 = s_pagedir_get_pml1(pml4, vptr);
    size_t pml1i = PML1I(vptr);
    if(!pml1[pml1i].p) return -ENOMAP;

    if(pml1[pml1i].mapping_type != JOVE_MAPPING_TYPE_UNTYPED) return -ENOMAP;
    uintptr_t untyped_pptr = (pml1[pml1i].pfn << 12);
    
    obj_settype(dst, JOVE_OBJ_UNTYPED);
    obj_setvalue(dst, untyped_pptr);

    jove_untyped_t *untyped_data = obj_deref(dst);
    untyped_data->limit = 0x1000;
    pagedir_invlpg(vptr);

    return 0;
}

int pagedir_unmap_shmem_frame(
        jove_owned_obj_t *tgt, uintptr_t vptr)
{
    page_t *pml4 = obj_deref(tgt);

    size_t pml4i = PML4I(vptr);
    if(!pml4[pml4i].p) return -ENOMAP;

    page_t *pml1 = s_pagedir_get_pml1(pml4, vptr);
    size_t pml1i = PML1I(vptr);
    if(!pml1[pml1i].p) return -ENOMAP;

    if(pml1[pml1i].mapping_type != JOVE_MAPPING_TYPE_UNTYPED) return -ENOMAP;

    pml1[pml1i].p = 0;
    pagedir_invlpg(vptr);

    return 0;
}

int pagedir_unmap_untyped_table(
        jove_owned_obj_t *tgt, jove_owned_obj_t *dst, uintptr_t vptr)
{
    page_t *target_table = NULL;
    page_t *pml4 = obj_deref(tgt);

    //If the PML4 entry does not exist at this virtual ptr, the user may be
    //trying to unmap the entire address space. We don't support this.
    size_t pml4i = PML4I(vptr);
    if(!pml4[pml4i].p) {
        return -ENOMAP;
    }
    target_table = &pml4[pml4i];

    //Try to unmap the pml4 entry if the pml3 entry does not exist.
    page_t *pml3 = (page_t*)(((uintptr_t)pml4[pml4i].pfn << 12) | IDENT_VIRT_BASE);
    size_t pml3i = PML3I(vptr);
    if(!pml3[pml3i].p) {
        goto unmap_table;
    }
    target_table = &pml3[pml3i];

    //Try to unmap the pml3 entry if the pml2 entry does not exist.
    page_t *pml2 = (page_t*)(((uintptr_t)pml3[pml3i].pfn << 12) | IDENT_VIRT_BASE);
    size_t pml2i = PML2I(vptr);
    if(!pml2[pml2i].p) {
        goto unmap_table;
    }
    //We are unmapping a pml2 entry.
    target_table = &pml2[pml2i];

unmap_table:
{
    uintptr_t frame = target_table->pfn << 12;
    page_t *target_vals = (page_t*)(frame | IDENT_VIRT_BASE);

    //Check that the page table at our target entry is empty.
    //If it's not, we return an error saying that the entry contains a map.
    for(size_t i = 0; i < 512; i++) {
        if(target_vals[i].p) return -EHASMAP;
    }

    //Populate the none object with an untyped object at the free'd page frame.
    obj_settype(dst, JOVE_OBJ_UNTYPED);
    obj_setvalue(dst, frame);

    jove_untyped_t *untyped = obj_deref(dst);
    untyped->limit = 0x1000;

    pagedir_invlpg(vptr);
    return 0;
}
}

void
paging_pml4_init(void)
{
    //Setup kernel pages
    memset(_kernel_pml4, 0, 0x1000);
    memset(_kernel_pml3, 0, 0x1000);
    memset(_kernel_pml2, 0, 0x1000);
    memset(_kernel_pml1, 0, 0x1000);

    memset(ident_pml3, 0, 0x1000 * IDENT_MASSIVEPAGES);
    memset(ident_pml2, 0, 0x1000 * IDENT_GIANTPAGES);

    /*Paging entries depend on physical addresses to properly map.
     * Since the virtual kernel address differs from the physical memory address,
     * we need to convert the pointers using the physical memory pointer passed
     * to us by the bootloader.*/
    uintptr_t kernel_pml4_pptr = KPTR_TOPHYS((uintptr_t)_kernel_pml4);
    uintptr_t kernel_pml3_pptr = KPTR_TOPHYS((uintptr_t)_kernel_pml3);
    uintptr_t kernel_pml2_pptr = KPTR_TOPHYS((uintptr_t)_kernel_pml2);
    uintptr_t kernel_pml1_pptr = KPTR_TOPHYS((uintptr_t)_kernel_pml1);

    /*Maps pml4-2 as global writeable pages.*/
    _kernel_pml4[511].value = kernel_pml3_pptr 
        | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL;
    _kernel_pml3[510].value = kernel_pml2_pptr 
        | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL;
    _kernel_pml2[0].value =   kernel_pml1_pptr 
        | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL;

    /*text*/
    extern void *_kernel_text_start;
    extern void *_kernel_text_end;
    uintptr_t kernel_text_start_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_text_start);
    uintptr_t kernel_text_end_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_text_end);
    size_t kernel_text_pages = ((kernel_text_end_pptr - kernel_text_start_pptr) / 0x1000) + 1;
    
    uintptr_t kernel_text_off = (((uintptr_t)&_kernel_text_start) - KERNEL_VIRT_BASE) >> 12;

    for(size_t i = 0; i < kernel_text_pages; i++) {
        size_t off = i + kernel_text_off;
        uintptr_t page_off = kernel_text_start_pptr + (i << 12);
        _kernel_pml1[off].value = page_off | PAGE_PRESENT | PAGE_GLOBAL;
    }

    /*rodata*/
    extern void *_kernel_rodata_start;
    extern void *_kernel_rodata_end;
    uintptr_t kernel_rodata_start_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_rodata_start);
    uintptr_t kernel_rodata_end_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_rodata_end);
    size_t kernel_rodata_pages = ((kernel_rodata_end_pptr - kernel_rodata_start_pptr) / 0x1000) + 1;

    uintptr_t kernel_rodata_off = kernel_text_off + kernel_text_pages;
    for(size_t i = 0; i < kernel_rodata_pages; i++) {
        size_t off = i + kernel_rodata_off;
        uintptr_t page_off = kernel_rodata_start_pptr + (i << 12);
        _kernel_pml1[off].value = page_off | PAGE_PRESENT | PAGE_GLOBAL | PAGE_XD;
    }

    /*data/bss*/
    extern void *_kernel_data_start;
    extern void *_kernel_bss_end;
    uintptr_t kernel_data_start_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_data_start);
    uintptr_t kernel_data_end_pptr = KPTR_TOPHYS((uintptr_t)&_kernel_bss_end);
    size_t kernel_data_pages = ((kernel_data_end_pptr - kernel_data_start_pptr) / 0x1000) + 1;

    uintptr_t kernel_data_off = kernel_rodata_off + kernel_rodata_pages;
    for(size_t i = 0; i < kernel_data_pages; i++) {
        size_t off = i + kernel_data_off;
        uintptr_t page_off = kernel_data_start_pptr + (i << 12);
        _kernel_pml1[off].value = page_off 
            | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL | PAGE_XD;
    }

    size_t kernel_stack_off = kernel_data_off + kernel_data_pages;
    __syscall_stack_guard = ((kernel_stack_off) << 12) + KERNEL_VIRT_BASE;
    __syscall_stack = __syscall_stack_guard + 0x1000;
    __interrupt_stack_guard = __syscall_stack + 0x1000;
    __interrupt_stack = __interrupt_stack_guard + 0x1000;

    __syscall_stack_mapping = (virtaddr_t*)&_kernel_pml1[kernel_stack_off + 1];
    __interrupt_stack_mapping = (virtaddr_t*)&_kernel_pml1[kernel_stack_off + 3];

    /*Identity pages are mapped as huge pages starting at the 256'th PML3 entry.
     *We want to support any range of physical memory up to several TiB, so
     *the mapper has to consider multiple heirarchies of huge pages.*/
    
    for(size_t l3 = 0; l3 < IDENT_MASSIVEPAGES; l3++) {
        uintptr_t ident_pml3_pptr = KPTR_TOPHYS(&ident_pml3[l3]);

        _kernel_pml4[256 + l3].value = ident_pml3_pptr 
            | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL;
    }

    for(size_t l2 = 0; l2 < IDENT_GIANTPAGES; l2++) {
        size_t l3 = l2 / 512;
        size_t l2_mod = l2 % 512;

        uintptr_t ident_pml2_pptr = KPTR_TOPHYS(&ident_pml2[l2]);
        ident_pml3[l3][l2_mod].value = ident_pml2_pptr 
            | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL;
    }

    for(size_t l1 = 0; l1 < PHYS_HUGEPAGES_MAX; l1++) {
        size_t l2 = l1 / 512;
        size_t l1_mod = l1 % 512;

        ident_pml2[l2][l1_mod].value = (l1 << 21) | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_PS | PAGE_GLOBAL | PAGE_XD;
        ident_pml2[l2][l1_mod].xd = 1;
    }

    paging_pml4_switch(kernel_pml4_pptr);

    {
        jove_owned_obj_t *pgdir_o = objdir_getlocal(_rootdir_obj, INIT_ROOTDIR_PAGEDIR_I);
        obj_settype(pgdir_o, JOVE_OBJ_PAGEDIR);
        obj_setvalue(pgdir_o, kernel_pml4_pptr);
        obj_release(pgdir_o);
    }

    /*Before tasking, if the kernel has a bug and crashes we need some
     * information about what happened.*/
    extern struct int_stackframe *(*__int_kernel_handles[])(struct int_stackframe*);
    __int_kernel_handles[IDT_EXCEPT_PF] = s_int_handle_pf_early;
}