#include "paging.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* 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;
}

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){
        PAGE_MAP(i);
    }

    // 4096 bytes stack
    PAGE_MAP_PHY(-4096, kvar_stack_pma);
    
    // Load new pml4
    u64 kpage_phy=((u64)kpages[0]-kvar_kernel_vma);
    asm volatile(
        "mov %0, %%rax        \n\t"
        "mov %%rax, %%cr3     \n\t"
        :: "r" (kpage_phy));
}

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 addess
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){
    u16 pml4=virt>>39&0x1FF;
    u16 pdp=virt>>30&0x1FF;
    u16 pd=virt>>21&0x1FF;
    u16 pt=virt>>12&0x1FF;
    options&=0xFFF; // Ensure options are on 12bits
    
    // Solve pdp
    if(pml4_table[pml4] == 0){
        pml4_table[pml4]=(u64)paging_allocate_table();
        pml4_table[pml4]|=options;
        paging_allocate_addr(pml4_table,virt,phy,options);
        return;
    }

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