#include "paging.h"
#include "arch/x86_64/x86_64.h"
#include "jove.h"
#include "init.h"
#include "config.h"
#include "hal/obj.h"
#include "lib/string.h"
#include "lib/printf.h"
#include "lib/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];

void*
phys_tovirt(uintptr_t phys)
{
    return (void*)(phys | 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));
}

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));
    kprintf("Faulting address: %p\n", cr2);
    kprintf("RIP: %p | RSP: %p\n", frame->int_rip, frame->int_rsp);
    print_state(&frame->state);

    hcf();
}

int pagedir_map_frame(
        jove_owned_obj_t *tgt, jove_owned_obj_t *untyped_obj, 
        uintptr_t vptr, struct syscall_pagedir_flags flg)
{
    page_t *pml4 = obj_deref(tgt);
    uintptr_t frame = obj_getvalue(untyped_obj);

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

    page_t *pml3 = (page_t*)(((uintptr_t)pml4[pml4i].paddr << 12) | IDENT_VIRT_BASE);
    size_t pml3i = PML3I(vptr);
    if(!pml3[pml3i].p) return ENOMAP;

    page_t *pml2 = (page_t*)(((uintptr_t)pml3[pml3i].paddr << 12) | IDENT_VIRT_BASE);
    size_t pml2i = PML2I(vptr);
    if(!pml2[pml2i].p) return ENOMAP;

    page_t *pml1 = (page_t*)(((uintptr_t)pml2[pml2i].paddr << 12) | IDENT_VIRT_BASE);
    size_t pml1i = PML1I(vptr);
    if(pml1[pml1i].p) return EEXIST;

    pml1[pml1i].value = frame |
        PAGE_PRESENT |
        PAGE_USER |
        (flg.writeable ? PAGE_WRITETHROUGH : 0) |
        (flg.executable ? 0 : PAGE_XD);
    return 0;
}

int pagedir_map_table(
        jove_owned_obj_t *tgt, jove_owned_obj_t *untyped_obj, 
        uintptr_t vptr, struct syscall_pagedir_flags flg)
{
    page_t *target_table = NULL;
    page_t *pml4 = obj_deref(tgt);
    uintptr_t frame = obj_getvalue(untyped_obj);

    void *frame_ident = obj_deref(untyped_obj);
    memset(frame_ident, 0, 0x1000);

    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].paddr << 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].paddr << 12) | IDENT_VIRT_BASE);
    size_t pml2i = PML2I(vptr);
    if(!pml2[pml2i].p) {
        target_table = &pml2[pml2i];
        goto map_table;
    }
    return EHASMAP;

map_table:
    target_table->value = frame |
        PAGE_PRESENT |
        PAGE_USER |
        (flg.writeable ? PAGE_WRITETHROUGH : 0) |
        (flg.executable ? 0 : PAGE_XD);
    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_text_start_pptr + (off << 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_text_start_pptr + (off << 12);
        _kernel_pml1[off].value = page_off | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_GLOBAL | PAGE_XD;
    }

    /*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, 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_handles[])(struct int_stackframe*);
    __int_handles[IDT_EXCEPT_PF] = s_int_handle_pf_early;
}