#include "paging.h"
#include "libc/stdio.h"
#include "libc/math.h"

/// Use a bitmap to keep track of allocated pages
char pages_status[PAGING_MAX_PAGES/8];
/// Kernel page directory (ATTENTION need to be 4096)
u32 k_pd[PAGING_MAX_DIR_ENTRY] __attribute__((aligned(4096)));
/// Kernel page table
u32 k_pt[PAGING_MAX_DIR_ENTRY][1024] __attribute__((aligned(4096)));

void paging_enable(){
    // Init pages status
    for(int i=0;i<PAGING_MAX_PAGES/8;i++)
        pages_status[i]=0;

    // Init page directory
    for(int i=0;i<PAGING_MAX_DIR_ENTRY;i++){
        k_pd[i]=0;
    }
    k_pd[0]=((int)&k_pt[0][0]);
    k_pd[0]|=7; // Permissions

    // Init page table 0
    int addr_offset=0;
    for(int i=0;i<1024;i++){
        k_pt[0][i]=addr_offset;
        k_pt[0][i]|=7; // Permission
        paging_set_usage(addr_offset,1); // Mark addr as used
        addr_offset+=PAGING_PAGE_SIZE; // 4Ko pages
    }

    // Allow access to more ram
    for(int i=1;i<PAGING_MAX_DIR_ENTRY;i++){
        k_pd[i]=((int)&k_pt[i][0]);
        k_pd[i]|=7; // Permissions
        for(int j=0;j<1024;j++){
            k_pt[i][j]=addr_offset;
            k_pt[i][j]|=7; // Permission
            addr_offset+=PAGING_PAGE_SIZE; // 4Ko pages
        }
    }

    // Turns on paging
    asm(
        "movl %0, %%eax       \n\t"
        "movl %%eax, %%cr3    \n\t" // Configure page table location
        "movl %%cr0, %%eax    \n\t"
        "orl %1, %%eax        \n\t"
        "movl %%eax, %%cr0    \n\t" // Turn on paging
        :: "b" (k_pd), "i" (PAGING_CR0_BIT)
    );
}
void paging_set_usage(int addr,char state){
    char bytes=pages_status[addr/PAGING_PAGE_SIZE/8];
    char bit=addr/PAGING_PAGE_SIZE%8;
    if(state=0)
        pages_status[addr/PAGING_PAGE_SIZE/8]=~(1<<bit)&bytes;
    else
        pages_status[addr/PAGING_PAGE_SIZE/8]=(1<<bit)|bytes;
}

void paging_dump(int min,int max){
    for(int i=0;i<(PAGING_MAX_PAGES/8);i++){
        for(int j=0;j<8;j++){
            char status=(pages_status[i]>>j)&0x1;
            if((i*8+j)>=min){
                if((i*8+j)<max || max<0)
                    printi(status);
            }
        }
    }
}

char* paging_allocate_next_page(){
    for(int i=0;i<(PAGING_MAX_PAGES/8);i++){
        char bytes=pages_status[i];
        for(int j=0;j<8;j++){
            char state=(bytes>>j)&1;
            if(state!=1){
                int page_id=i*8+j;
                int page_addr=PAGING_PAGE_SIZE*page_id;
                paging_set_usage(page_addr,1);
                return((char*)page_addr);
            }
        }
    }
    print("Could not allocate anymore pages! Stopping...");
    asm("hlt");
}

int *paging_allocate(int p){
    // ----- Populate kernel adresses (to be able to execute interrupts codes)
    int *page_dir=(int*)paging_allocate_next_page();
    int *k_page_table=(int*)paging_allocate_next_page();
    // Kernel is located in the first 4Mb (page dir entry 0)
    page_dir[0]=(int)k_page_table|3;
    int addr_offset=0;
    for(int i=0;i<1024;i++){
        k_page_table[i]=addr_offset;
        k_page_table[i]|=3; // Permissions
        addr_offset+=PAGING_PAGE_SIZE; // 4Ko pages
    }
  
    // ----- Populate task adresses
    int *u_page_table=(int)paging_allocate_next_page();
    page_dir[1]=(int)u_page_table|7;
    int dir_entry=1; // Entry point is at 1024*4096=4194304 (page dir entry 1 and page table entry 0)
    int pt_entry=0;
    int p_current=max(1,p); // Allocate at least 1 page
    p_current++; // Allocated one more page for the kernel stack (one per task)
    while(p_current!=0){
        u_page_table[pt_entry]=(int)paging_allocate_next_page();
        if(p_current==1) // Kernel stack is the last page
            u_page_table[pt_entry]|=3;
        else
            u_page_table[pt_entry]|=7;
        p_current--;
        pt_entry++;
        if(pt_entry%1024==0){
            dir_entry++;
            pt_entry=0;
            u_page_table=(int*)paging_allocate_next_page();
            page_dir[dir_entry]=(int)u_page_table;
            if(p_current==1) // Kernel stack is the last page (and might be in the last page directory entry)
                 page_dir[dir_entry]|=3;
            else
                 page_dir[dir_entry]|=7;
        }
    }
    return page_dir;
}

void paging_page_fault(){
  print("Page fault!");
  asm("hlt");
}