#include <lock.h>
#include <stdatomic.h>
#include <stdint.h>
#include <kprint.h>
#include <neobbo.h>
#include <string.h>
#include <limine.h>
#include <mm/pmm.h>
#include <mm/vmm.h>
#include <sys/acpi.h>
#include <arch/amd64/hal/apic.h>

struct limine_executable_address_request kernel_addr_request = {
    .id = LIMINE_EXECUTABLE_ADDRESS_REQUEST_ID,
    .revision = 0
};

struct limine_executable_address_response *kernel_address;

extern uint64_t hhdmoffset;

uint64_t *kernel_page_map = 0;
uint64_t kernel_virt = 0;

uint64_t text_start, text_end, rodata_start, rodata_end, data_start, data_end;

uint64_t kernel_start, kernel_end;

void vmm_set_ctx(uint64_t *page_map){
    __asm__ volatile (
        "movq %0, %%cr3\n"
        : : "r" ((uint64_t)((uint64_t)(page_map) - get_kinfo()->hhdmoffset)) : "memory"
    );
}


void vmm_init(){
    
    struct limine_executable_address_response *kernel_address = kernel_addr_request.response;

    if(!kernel_address){
        klog(__func__, "Kernel address not recieved");
    }

    kernel_page_map = (uint64_t*)((uint64_t)pmm_alloc() + get_kinfo()->hhdmoffset);

    if(!kernel_page_map){
        klog(__func__, "Allocating block for page map failed");
    }

    memset(kernel_page_map, 0, PAGE_SIZE);

    // map kernel, stolen
    extern link_symbol_ptr text_start_addr, text_end_addr,
        rodata_start_addr, rodata_end_addr,
        data_start_addr, data_end_addr;

        text_start = ALIGN_DOWN((uint64_t)text_start_addr, PAGE_SIZE),
        rodata_start = ALIGN_DOWN((uint64_t)rodata_start_addr, PAGE_SIZE),
        data_start = ALIGN_DOWN((uint64_t)data_start_addr, PAGE_SIZE),
        text_end = ALIGN_UP((uint64_t)text_end_addr, PAGE_SIZE),
        rodata_end = ALIGN_UP((uint64_t)rodata_end_addr, PAGE_SIZE),
        data_end = ALIGN_UP((uint64_t)data_end_addr, PAGE_SIZE); 


    // map usable entries, framebuffer and bootloader reclaimable shit
    extern struct limine_memmap_response *memmap_response;
    for(uint64_t i = 0; i < memmap_response->entry_count; i++){
        if(memmap_response->entries[i]->type == LIMINE_MEMMAP_USABLE){
            for(uint64_t j = 0; j < memmap_response->entries[i]->length; j+=PAGE_SIZE){
                vmm_map_page(kernel_page_map, memmap_response->entries[i]->base+j+get_kinfo()->hhdmoffset, memmap_response->entries[i]->base+j, PTE_BIT_PRESENT | PTE_BIT_RW);
            }
        }
        if(memmap_response->entries[i]->type == LIMINE_MEMMAP_FRAMEBUFFER){
            for(uint64_t j = 0; j < memmap_response->entries[i]->length; j+=PAGE_SIZE){
                vmm_map_page(kernel_page_map, memmap_response->entries[i]->base+j+get_kinfo()->hhdmoffset, memmap_response->entries[i]->base+j, PTE_BIT_PRESENT | PTE_BIT_RW | PTE_BIT_NX);
            }
        }
        if(memmap_response->entries[i]->type == LIMINE_MEMMAP_BOOTLOADER_RECLAIMABLE){
            for(uint64_t j = 0; j < memmap_response->entries[i]->length; j+=PAGE_SIZE){
                vmm_map_page(kernel_page_map, memmap_response->entries[i]->base+j+get_kinfo()->hhdmoffset, memmap_response->entries[i]->base+j, PTE_BIT_PRESENT | PTE_BIT_RW);
            }
        }
        if(memmap_response->entries[i]->type == LIMINE_MEMMAP_ACPI_RECLAIMABLE){
            for(uint64_t j = 0; j < memmap_response->entries[i]->length; j+=PAGE_SIZE){
                vmm_map_page(kernel_page_map, memmap_response->entries[i]->base+j+get_kinfo()->hhdmoffset, memmap_response->entries[i]->base+j, PTE_BIT_PRESENT | PTE_BIT_RW);
            }
        }

    }

    for (uintptr_t text_addr = text_start; text_addr < text_end; text_addr += PAGE_SIZE) {
        uintptr_t phys = text_addr - kernel_address->virtual_base + kernel_address->physical_base;
        vmm_map_page(kernel_page_map, text_addr, phys, PTE_BIT_PRESENT);
    }
    /* Kernel starts with the text section */
    kernel_start = text_start;

    kprintf("vmm: text_start: 0x{xn}vmm: text_end: 0x{xn}", text_start, text_end);

    for (uintptr_t rodata_addr = rodata_start; rodata_addr < rodata_end; rodata_addr += PAGE_SIZE) {
        uintptr_t phys = rodata_addr - kernel_address->virtual_base + kernel_address->physical_base;
        vmm_map_page(kernel_page_map, rodata_addr, phys, PTE_BIT_PRESENT | PTE_BIT_NX);
    }

    kprintf("vmm: rodata_start: 0x{xn}vmm: rodata_end: 0x{xn}", rodata_start, rodata_end);

    for (uintptr_t data_addr = data_start; data_addr < data_end; data_addr += PAGE_SIZE) {
        uintptr_t phys = data_addr - kernel_address->virtual_base + kernel_address->physical_base;
        vmm_map_page(kernel_page_map, data_addr, phys, PTE_BIT_PRESENT | PTE_BIT_RW | PTE_BIT_NX);
    }

    kprintf("vmm: data_start: 0x{xn}vmm: data_end: 0x{xn}", data_start, data_end);

    /* Kernel ends with the data section */
    kernel_end = data_end;

    vmm_set_ctx(kernel_page_map);

    asm volatile(
        "movq %%cr3, %%rax\n\
	   movq %%rax, %%cr3\n"
        : : : "rax"
   );
    
}

uint64_t *get_lower_table(uint64_t *page_map, uint64_t offset){

    if((page_map[offset] & PTE_BIT_PRESENT) != 0){
        return (uint64_t*)( ((uint64_t)page_map[offset] & 0x000ffffffffff000) + get_kinfo()->hhdmoffset);
    }


    uint64_t *ret = pmm_alloc();

    if(!ret){
        klog(__func__, "Failed to allocate page table");
        kprintf("page_map: 0x{xn}", (uint64_t)page_map);
        kprintf("offset: 0x{xn}", offset);
        
        return NULL;
    }

    memset((uint64_t*)((uint64_t)ret + get_kinfo()->hhdmoffset), 0, PAGE_SIZE);

    page_map[offset] = (uint64_t)ret | PTE_BIT_PRESENT | PTE_BIT_RW | PTE_BIT_US;

    return (uint64_t*)((uint64_t)ret + get_kinfo()->hhdmoffset);


}

atomic_flag page_table_lock = ATOMIC_FLAG_INIT;

void vmm_map_page(uint64_t *page_map, uint64_t virt_addr, uint64_t phys_addr, uint64_t flags){
    /* Probably slow, fix in future */
    acquire_spinlock(&page_table_lock);

    uint64_t pml4_offset = (virt_addr >> 39) & 0x1ff;
    uint64_t pdp_offset = (virt_addr >> 30) & 0x1ff;
    uint64_t pd_offset = (virt_addr >> 21) & 0x1ff;
    uint64_t pt_offset = (virt_addr >> 12) & 0x1ff;

    uint64_t *pdp = get_lower_table(page_map, pml4_offset);

    if(!pdp){
        klog( __func__, "Failed to allocate PDP");
        kkill();
    }

    uint64_t *pd = get_lower_table(pdp, pdp_offset);

    if(!pd){
        klog(__func__, "Failed to allocate PD");
        kkill();
    }

    uint64_t *pt = get_lower_table(pd, pd_offset);

    if(!pt){
        klog( __func__, "Failed to allocate PT");
        kkill();
    }

    if(pt[pt_offset] != 0){
        goto end;
    }

    pt[pt_offset] = phys_addr | flags;

    asm volatile(
        "movq %%cr3, %%rax\n\
	   movq %%rax, %%cr3\n"
        : : : "rax"
   );

   end:

   free_spinlock(&page_table_lock);

}

void vmm_free_page(uint64_t *page_map, uint64_t virt_addr){
    uint64_t pml4_offset = (virt_addr >> 39) & 0x1ff;
    uint64_t pdp_offset = (virt_addr >> 30) & 0x1ff;
    uint64_t pd_offset = (virt_addr >> 21) & 0x1ff;
    uint64_t pt_offset = (virt_addr >> 12) & 0x1ff;

    uint64_t *pdp = get_lower_table(page_map, pml4_offset);


    if(!pdp){
        klog(__func__, "Failed to allocate PDP");
        kkill();
    }

    uint64_t *pd = get_lower_table(pdp, pdp_offset);

    if(!pd){
        klog( __func__, "Failed to allocate PD");
        kkill();
    }


    uint64_t *pt = get_lower_table(pd, pd_offset);

    if(!pt){
        klog(__func__, "Failed to allocate PT");
        kkill();
    }

    /* Free the page at the physical address pointed by the pt entry */
    pmm_free((uint64_t*)(pt[pt_offset] & 0x000ffffffffff000));

    /* Set it to zero (mark as not present) */
    pt[pt_offset] = 0;

    asm volatile(
        "movq %%cr3, %%rax\n\
	    movq %%rax, %%cr3\n"
        : : : "rax"
   );
}

uint64_t vmm_get_phys_addr(uint64_t *page_map, uint64_t virt_addr){
    uint64_t pml4_offset = (virt_addr >> 39) & 0x1ff;
    uint64_t pdp_offset = (virt_addr >> 30) & 0x1ff;
    uint64_t pd_offset = (virt_addr >> 21) & 0x1ff;
    uint64_t pt_offset = (virt_addr >> 12) & 0x1ff;

    uint64_t pml4e = page_map[pml4_offset];
    if (!(pml4e & 1)) return 0; // 
    uint64_t *pdp = (uint64_t *)((pml4e & 0x000ffffffffff000) + get_kinfo()->hhdmoffset);

    uint64_t pdpe = pdp[pdp_offset];
    if (!(pdpe & 1)) return 0; 
    uint64_t *pd = (uint64_t *)((pdpe & 0x000ffffffffff000) + get_kinfo()->hhdmoffset);

    uint64_t pde = pd[pd_offset];
    if (!(pde & 1)) return 0; 
    uint64_t *pt = (uint64_t *)((pde & 0x000ffffffffff000) + get_kinfo()->hhdmoffset);

    uint64_t pte = pt[pt_offset];
    if (!(pte & 1)) return 0;

    return pte & 0x000ffffffffff000;
}

uint64_t kget_phys_addr(uint64_t *virt_addr){
    return vmm_get_phys_addr(kernel_page_map, (uint64_t)virt_addr);
}


/* Maps `size` number of free pages at the specified virtual address */
int vmm_map_contigious_pages(uint64_t *page_map, uint64_t virt_addr, uint64_t phys_addr, uint64_t size, uint64_t flags){
    for(uint64_t i = 0; i < size; i++){
        vmm_map_page(page_map, virt_addr + i * PAGE_SIZE, (uint64_t)phys_addr, flags);
    }

    return 0;
}

#define VA_BASE 0x900915000
uint64_t va_base = VA_BASE;
atomic_flag va_lock = ATOMIC_FLAG_INIT;

/* Allocates some pages from the PMM and makes them contigious in the kernels memory map. */
void *va_alloc_contigious_pages(size_t pages){
    acquire_spinlock(&va_lock);

    if(pages == 0){
        free_spinlock(&va_lock);
        return NULL;
    }

    size_t i;
    for(i = 0; i < pages; i++){
        void *t = pmm_alloc();

        if(t == NULL){
            free_spinlock(&va_lock);
            return NULL;
        }

        vmm_map_page(kernel_page_map, va_base+get_kinfo()->hhdmoffset+i*PAGE_SIZE, (uint64_t)t, PTE_BIT_RW | PTE_BIT_PRESENT);
    }
    
    uint64_t va_base_old = va_base;

    va_base += i*PAGE_SIZE;

    free_spinlock(&va_lock);

    return (void*)(va_base_old+get_kinfo()->hhdmoffset);

}

/* Maps pages from phys_addr to phys_addr+size into the kernels address space */
void kmap_pages(void *phys_addr, uint64_t size, uint64_t flags){
    for(uint64_t i = 0; i < size; i++){
        vmm_map_page(kernel_page_map, (uint64_t)phys_addr + get_kinfo()->hhdmoffset + (i * PAGE_SIZE), (uint64_t)phys_addr + (i * PAGE_SIZE), PTE_BIT_PRESENT | flags);
    }
}

void kunmap_pages(void *addr, uint64_t size){
    for(uint64_t i = 0; i < size; i++){
        vmm_free_page(kernel_page_map, (uint64_t)addr + i*PAGE_SIZE);
    }
}
