#include "paging.hpp"
#include "core/asm.hpp"
#include "core/types.hpp"
#include "libs/stdio.hpp"
#include "libs/string.hpp"

char paging_status[PAGING_MAX_PAGE / 8];
u64 kpages[MAX_TABLES][512] __attribute__((aligned(4096)));
int kpages_next=1; // First page is for the pml4
u64* kpml4;

u64* paging_allocate_table(){
    u64 addr=(u64)kpages[kpages_next];
    u64* allocated=(u64*)(addr-kvar_kernel_vma);
    kpages_next++;
    if(kpages_next>=PAGING_MAX_PAGE){
        printk("Could not allocate more page structures. Kernel Panic!");
        while(1);
    }
    return allocated;
}
u64* paging_allocate_utable(){
    u64 *table=PAGE_ALLOCATE();
    for(u32 i=0;i<512;i++)
        table[i]=0;
    return table;
}

void paging_enable() {
    // Init status
    for (int i = 0; i < PAGING_MAX_PAGE / 8; i++) {
        paging_status[i] = 0;
    }

    // Init tables
    for(int i=0;i<MAX_TABLES;i++){
        for(int j=0;j<512;j++)
            kpages[i][j]=0;
    }

    // Allocate paging for the kernel (to not override the source
    // code during the next paging_allocate_table() calls)
    paging_allocate_contiguous(kvar_userspace_pma/4096);

    // Setting up new kernel address space
    for(u64 i=0;i<=0x10000000;i+=4096){
        // Higher half mapping
        PAGE_VIRT_MAP(i,PAGING_OPT_DEFAULTS);
        // Allow access to RAM:
        paging_allocate_addr(kpages[0], i, i, PAGING_OPT_DEFAULTS, 1);
    }

    // 4096 bytes stack
    PAGE_MAP(kvar_kernel_vma-4096, kvar_stack_pma,PAGING_OPT_DEFAULTS);

    // Load new pml4
    kpml4=(u64*)((u64)kpages[0]-kvar_kernel_vma);
    lpml4(kpml4);
}

u64* paging_allocate_contiguous(int npages){
    int n_contiguous=0;
    for (int i = 0; i < PAGING_MAX_PAGE / 8; i++) {
        for (int j = 0; j < 8; j++) {
            char bit=(paging_status[i]&(0x1<<j))>>j;
            if(bit!=1){
                n_contiguous++;
            }
            else {
                n_contiguous=0;
            }
            if(n_contiguous==npages){
                n_contiguous--; // Since we use it now as index, not a counter
                int start_page=(i*8+j)-n_contiguous;
                while(n_contiguous>=0){
                    int cur_page=(i*8+j)-n_contiguous;
                    paging_status[cur_page/8]|=(0x1<<(cur_page%8)); // Allocate
                    n_contiguous--;
                }
                u64 phy_addr=(4096*start_page);
                return (u64*) phy_addr;
            }
        }   
    }

    printk("Could not allocate %d contiguous pages. Kernel panic!",npages);
    while(1);
    return 0;
}

void paging_deallocate(u64 addr){
    u64 page_number=PAGE(addr)/4096;
    char byte=paging_status[page_number/8];
    paging_status[page_number/8]=byte&(~(1<<(page_number%8)));
}


/// TODO: Debug address
void paging_deallocate_pml4(u64* pml4){
    for(int i=0;i<512;i++){
        u64* pdp=(u64*)PAGE(pml4[i]);
        if(pml4[i]==0)
            continue;
        for(int j=0;j<512;j++){
            u64* pd=(u64*)PAGE(pdp[j]);
            if(pdp[j]==0)
                continue;
            for(int k=0;k<512;k++){
                u64* pt=(u64*)PAGE(pd[k]);
                if(pd[k]==0)
                    continue;
                for(int l=0;l<512;l++){
                    if(pt[l]==0)
                        continue;
                    paging_deallocate_table((u64*)PAGE(pt[l]));
                }
                paging_deallocate_table((u64*)PAGE(pd[k]));
            }
            paging_deallocate_table((u64*)PAGE(pdp[j]));
        }
        paging_deallocate_table((u64*)PAGE(pml4[i]));
    }
    paging_deallocate_table((u64*)PAGE((u64)pml4));
}

void paging_dump(int min, int max) {
    for (int i = 0; i < PAGING_MAX_PAGE / 8; i++) {
        if(i>=min && i<=max){
            printk("Byte %d   ", i);
            for (int j = 0; j < 8; j++) {
                char bit = (paging_status[i] & (0x1 << j)) >> j;
                printk("%d", bit);
            }
            print("\n");
        }
    }
}

void paging_deallocate_table(u64* table){
    char *c_table=(char*)PAGE((u64)table);
    for(u8 i=0;i<8;i++){
        paging_deallocate((u64)c_table);
        c_table+=4096;
    }
}

void paging_allocate_addr(u64* pml4_table, u64 virt, u64 phy, u16 options, char useKernelTables){
    u16 pml4=virt>>39&0x1FF;
    u16 pdp=virt>>30&0x1FF;
    u16 pd=virt>>21&0x1FF;
    u16 pt=virt>>12&0x1FF;
    options&=0xFFF; // Ensure 12 bits
    
    // Solve pdp
    if(pml4_table[pml4] == 0){
        pml4_table[pml4]=(u64)(useKernelTables ? paging_allocate_table() : paging_allocate_utable());
        pml4_table[pml4]|=options;
        paging_allocate_addr(pml4_table,virt,phy,options,useKernelTables);
        return;
    }
    
    // Solve pd
    u64* pdp_table=(u64*)(PAGE(pml4_table[pml4]));
    pdp_table=useKernelTables ? VIRT(pdp_table) : pdp_table;
    if(pdp_table[pdp] == 0){
        pdp_table[pdp]=(u64)(useKernelTables ? paging_allocate_table() : paging_allocate_utable());
        pdp_table[pdp]|=options;
        paging_allocate_addr(pml4_table,virt,phy,options,useKernelTables);
        return;
    }

    // Solve pt
    u64* pd_table=(u64*)(PAGE(pdp_table[pdp]));
    pd_table=useKernelTables ? VIRT(pd_table) : pd_table;
    if(pd_table[pd] == 0){
        pd_table[pd]=(u64)(useKernelTables ? paging_allocate_table() : paging_allocate_utable());
        pd_table[pd]|=options;
        paging_allocate_addr(pml4_table,virt,phy,options,useKernelTables);
        return;
    }

    // Solve address
    u64* pt_table=(u64*)(PAGE(pd_table[pd]));
    pt_table=useKernelTables ? VIRT(pt_table) : pt_table;
    pt_table[pt]=PAGE(phy);
    pt_table[pt]|=options;
    return;
}

u64* paging_create_task(int npages){
    u64 *pml4=paging_allocate_utable();

    int i;
    for(i=0;i<npages;i++){
        paging_allocate_addr(pml4, i*4096, (u64)PAGE_ALLOCATE(), PAGING_OPT_DEFAULTS|PAGING_OPT_US, 0);
    }
    // Allocate a page for the user stack
    paging_allocate_addr(pml4, i*4096, (u64)PAGE_ALLOCATE(), PAGING_OPT_DEFAULTS|PAGING_OPT_US, 0);
    // Allocate a page for the kernel stack
    paging_allocate_addr(pml4, (i+1)*4096, (u64)PAGE_ALLOCATE(), PAGING_OPT_DEFAULTS, 0);

    // Enable kernel access
    u16 pml4_entry=kvar_kernel_vma>>39&0x1FF;
    pml4[pml4_entry]=kpages[0][pml4_entry];

    return pml4;
}