/**
 * The usage of data that could goes
 * to the .bss section is forbidden here
 */

#include "core/types.hpp"
#include "core/paging.hpp"


/// @brief Define where first PAE paging paging will be
u64 trampoline_next_page;

/// @brief Ensure
#define MAX_MB_INFOS 4096
u8 mb_infos[MAX_MB_INFOS];

u64 trampoline_paging_allocate_table(){
    for(u64 i=0;i<512;i++)
        ((u64*)trampoline_next_page)[i]=0;
    u64 ret=trampoline_next_page;
    trampoline_next_page+=4096;
    return ret;
}


void trampoline_paging_allocate_addr(u64* pml4_table, u64 virt, u64 phy, u16 options){
    u64 pml4=virt>>39&0x1FF;
    u64 pdp=virt>>30&0x1FF;
    u64 pd=virt>>21&0x1FF;
    u64 pt=virt>>12&0x1FF;
    options&=0xFFF; // Ensure options are on 12bits
    

    // Solve pdp
    if(pml4_table[pml4] == 0){
        pml4_table[pml4]=(u64)trampoline_paging_allocate_table();
        pml4_table[pml4]|=options;
        trampoline_paging_allocate_addr(pml4_table,virt,phy,options);
        return;
    }

    // Solve pd
    u64* pdp_table=(u64*)PAGE(pml4_table[pml4]);
    if(pdp_table[pdp] == 0){
        pdp_table[pdp]=(u64)trampoline_paging_allocate_table();
        pdp_table[pdp]|=options;
        trampoline_paging_allocate_addr(pml4_table,virt,phy,options);
        return;
    }
    // Solve pt
    u64* pd_table=(u64*)PAGE(pdp_table[pdp]);
    if(pd_table[pd] == 0){
        pd_table[pd]=(u64)trampoline_paging_allocate_table();
        pd_table[pd]|=options;
        trampoline_paging_allocate_addr(pml4_table,virt,phy,options);
        return;
    }
    // Solve address
    u64* pt_table=(u64*)PAGE(pd_table[pd]);
    if(pt_table[pt] == 0){
        pt_table[pt]=PAGE(phy);
        pt_table[pt]|=options;
        return;
    }
}

/**
 * Setup High-Half Kernel Paging
 */
extern "C" u8* trampoline(u8* mb_infos_addr){
    // Save mb infos before anything else
    for(int i=0;i<MAX_MB_INFOS;i++){
        mb_infos[i]=mb_infos_addr[i];
    }

    u64 kernel_vma,stack_pma;
    asm("movq $__kernel_vma, %0":"=r"(kernel_vma));
    asm("movq $__userspace_pma, %0":"=r"(trampoline_next_page));
    asm("movq $__stack_pma, %0":"=r"(stack_pma));

    // ----- Build the trampoline paging tables
    u64 *pml4=(u64*)trampoline_paging_allocate_table();
    // Higher Half kernel
    for(u64 i=0;i<=0x1000000;i+=4096){
        trampoline_paging_allocate_addr(pml4,kernel_vma+i,i,0x3);
    }
    // Stack
    trampoline_paging_allocate_addr(pml4,kernel_vma-4096,stack_pma,0x3);
    // First bytes for current eip
    for(u64 i=0;i<=0x1000000;i+=4096){
        trampoline_paging_allocate_addr(pml4,i,i,0x3);
    }

    // Load new pml4
    asm volatile(
        "movq %0, %%rax        \n\t"
        "movq %%rax, %%cr3     \n\t"
        :: "r" (pml4));
        
    return mb_infos;
}