#include "pagetable.h"
#include "pageframealloc.h"
#include "../utils/string.h"

void load_cr3( void* cr3_value ) {
    asm volatile("mov %0, %%cr3" :: "r"((uint64_t)cr3_value) : "memory");
}

void write_identity_map(pagetable_t* identity_table, uint64_t size){
    for(uintptr_t i=0; i<size/4096; i++){
        identity_table[i/512].entries[i%512] = {
            .present = 1,
            .write = 1,
            .user = 0,
            .address = i,
        };
    }
}

void create_basic_page_table(four_lvl_paging_t* table, pagetable_t* kernel_identity){

    assert(table);
    assert(kernel_identity);

    table->l4 = (pagetable_t*)PageFrameAllocator::alloc(true);
    table->l3 = (pagetable_t*)PageFrameAllocator::alloc(true);
    table->l2 = (pagetable_t*)PageFrameAllocator::alloc(true);

    //memset(table, 0, sizeof(four_lvl_paging_t));
    table->l4->entries[0] = {
        .present = 1,
        .write = 1,
        .user = 1,
        .address = ((uintptr_t)(table->l3)) >> 12,
    };

    table->l3->entries[0] = {
        .present = 1,
        .write = 1,
        .user = 1,
        .address = ((uintptr_t)(table->l2)) >> 12,
    };

    //kernel memory
    for(uint32_t i=0; i<KERNEL_MEMORY_BORDER/4096/512; i++){
        table->l2->entries[i] = {
            .present = 1,
            .write = 1,
            .user = 1,
            .address = (uintptr_t)(&kernel_identity[i]) >> 12,
        };
    }

    //write_identity_map(table->l1, KERNEL_MEMORY_BORDER);
    //for(uintptr_t i=0; i<KERNEL_MEMORY_BORDER/4096; i++){
    //    table->l1[i/512].entries[i%512] = {
    //        .present = 1,
    //        .write = 1,
    //        .user = 0,
    //        .address = i,
    //    };
    //}

    //add ioapic address to pagetable
    table->l2_apic = (pagetable_t*)PageFrameAllocator::alloc(true);
    table->ioapic = (pagetable_t*)PageFrameAllocator::alloc(true);
    table->lapic = (pagetable_t*)PageFrameAllocator::alloc(true);

    table->l3->entries[3] = {
        .present = 1,
        .write = 1,
        .user = 1,
        .address = (uintptr_t)(table->l2_apic) >> 12,
    };
    table->l2_apic->entries[502] = {
        .present = 1,
        .write = 1,
        .user = 1,
        .address = (uintptr_t)(table->ioapic) >> 12,
    };
    table->ioapic->entries[0] = {
        .present = 1,
        .write = 1,
        .user = 0,
        .address = 0xFEC00,
    };

    //lapic
    table->l2_apic->entries[503] = {
        .present = 1,
        .write = 1,
        .user = 1,
        .address = (uintptr_t)(table->lapic) >> 12,
    };
    table->lapic->entries[0] = {
        .present = 1,
        .write = 1,
        .user = 0,
        .address = 0xFEE00,
    };

    //unmap 0 page
    kernel_identity[0].entries[0].present = 0;
}

uintptr_t isMapped(uintptr_t vaddr, four_lvl_paging_t* flpt){
    uint16_t l4Index = (vaddr>>39) & 0x1FF;
    uint16_t l3Index = (vaddr>>30) & 0x1FF;
    uint16_t l2Index = (vaddr>>21) & 0x1FF;
    uint16_t l1Index = (vaddr>>12) & 0x1FF;

    if(flpt->l4->entries[l4Index].present){
        pagetable_t* lvl3 = (pagetable_t*)(flpt->l4->entries[l4Index].address<<12);
        if(lvl3->entries[l3Index].present){
            pagetable_t* lvl2 = (pagetable_t*)(lvl3->entries[l3Index].address<<12);
            if(lvl2->entries[l2Index].present){
                pagetable_t* lvl1 = (pagetable_t*)(lvl2->entries[l2Index].address<<12);
                if(lvl1->entries[l1Index].present)
                    return lvl1->entries[l1Index].address<<12;
            }
        }
    }
    return 0;
}

void* getFreeVirtSpace(four_lvl_paging_t* search_table, uint8_t num_pages){
    uint32_t start_v = 0x4000;
    uint32_t stop_v = 0x6000;
    static uint32_t next_start_v = 0x4000;
    //static uint32_t next_start_v = start_v;

    //start from last found address
    for (uint32_t v=next_start_v; v<stop_v; v++) {
        bool space_is_free = true;
        for(uint32_t i=0; i<num_pages; i++){
            if(isMapped((uintptr_t)(v+i)<<12, search_table)) {
                space_is_free = false;
            }
        }
        if(space_is_free){
            //next_start_v = v+num_pages;
            return  (void*) ((uintptr_t) (v <<12));
        }
    }
    //start again from the real start address
    for (uint32_t v=start_v; v<next_start_v; v++) {
        bool space_is_free = true;
        for(uint32_t i=0; i<num_pages; i++){
            if(isMapped((uintptr_t)(v+i)<<12, search_table)) {
                space_is_free = false;
            }
        }
        if(space_is_free){
            next_start_v = v+num_pages;
            return  (void*) ((uintptr_t) (v <<12));
        }
    }
    return 0;
}


void setMapping(uintptr_t vaddr, void* frame, four_lvl_paging_t* flpt, bool write){
    uint16_t l4Index = (vaddr>>39) & 0x1FF;
    uint16_t l3Index = (vaddr>>30) & 0x1FF;
    uint16_t l2Index = (vaddr>>21) & 0x1FF;
    uint16_t l1Index = (vaddr>>12) & 0x1FF;

    if(!(flpt->l4->entries[l4Index].present)){
        pagetable_t* newl3 = (pagetable_t*)PageFrameAllocator::alloc(true);
        memset(newl3, 0, 4096);
        flpt->l4->entries[l4Index] = {
            .present = 1,
            .write = 1,
            .user = 1,
            .address = (uintptr_t)newl3 >> 12
        };
    }

    pagetable_t* lvl3 = (pagetable_t*)(flpt->l4->entries[l4Index].address<<12);
    if(!(lvl3->entries[l3Index].present)){
        pagetable_t* newl2 = (pagetable_t*)PageFrameAllocator::alloc(true);
        memset(newl2, 0, 4096);
        lvl3->entries[l3Index] = {
            .present = 1,
            .write = 1,
            .user = 1,
            .address = (uintptr_t)newl2 >> 12
        };
    }

    pagetable_t* lvl2 = (pagetable_t*)(lvl3->entries[l3Index].address<<12);
    if(!(lvl2->entries[l2Index].present)){
        pagetable_t* newl1 = (pagetable_t*)PageFrameAllocator::alloc(true);
        memset(newl1, 0, 4096);
        lvl2->entries[l2Index] = {
            .present = 1,
            .write = write,
            .user = 1,
            .address = (uintptr_t)newl1 >> 12
        };
    }

    pagetable_t* lvl1 = (pagetable_t*)(lvl2->entries[l2Index].address<<12);

    if(frame){
        assert(!(lvl1->entries[l1Index].present)); // should not be present, bc its a new mapping
        lvl1->entries[l1Index] = {
            .present = 1,
            .write = write,
            .user = 1,
            .address = (uintptr_t)frame >> 12
        };
    }
    else //unmap if nullptr
        lvl1->entries[l1Index].present = 0;
}

void copy_pagetable(four_lvl_paging_t* parent_table, four_lvl_paging_t* child_table){
    //can safely assume only one lvl4 entry
    pagetable_t* lvl3 = (pagetable_t*)(parent_table->l4->entries[0].address<<12);
    for(uint16_t i3=0; i3<512; i3++){
        if(lvl3->entries[i3].present){
            pagetable_t* lvl2 = (pagetable_t*)(lvl3->entries[i3].address<<12);
            for(uint16_t i2=0; i2<512; i2++){
                if(lvl2->entries[i2].present){
                    pagetable_t* lvl1 = (pagetable_t*)(lvl2->entries[i2].address<<12);
                    for(uint16_t i1=0; i1<512; i1++){
                        if(lvl1->entries[i1].present){
                            uintptr_t vaddr =
                                ((uintptr_t)i3<<30) |
                                ((uintptr_t)i2<<21) |
                                ((uintptr_t)i1<<12) ;
                            if(vaddr < 0x4000000)
                                continue;
                            if(!isMapped(vaddr, child_table)){
                                //only copy if not part of basic pagetable
                                assert( (i3>0) || (i2>=32)); //assert user memory 
                                void* frame = PageFrameAllocator::alloc(false);
                                setMapping(vaddr, frame, child_table);

                                //map to local space for copy
                                void* dest = getFreeVirtSpace(parent_table, 1);
                                setMapping((uintptr_t)dest, frame, parent_table);
                                memcpy(dest, (void*)vaddr, 4096);
                                setMapping((uintptr_t)dest, 0, parent_table);
                            }
                        }
                    }
                }
            }
        }
    }
}