#include "paging.h"
#include "pmm.h"
#include <kernel.h>
#include <stddef.h>
#include <limine.h>

/*
Paging on x86 uses four different page table levels:
cr3 register contains the phys address for the PML4 (root directory)

Each directory/table is made of 512 entries, each one uint64_t
Each of these entries have special bits (PRESENT/WRITEABLE/USER/etc.)
that dictates their attributes. Also these bits fall back on children tables.

If we use 1GB huge pages: PML4 -> PDPT -> 1gb pages
2MB huge pages: PML4 -> PDPT -> PD -> 2mb pages
4KB (regular size): PML4 -> PDPT -> PD -> PT -> 4kb pages
*/

static inline void load_cr3(uint64_t value) {
    asm volatile ("mov %0, %%cr3" :: "r"(value) : "memory");
}

// To flush TLB
static inline void invlpg(void *addr)
{
    asm volatile("invlpg (%0)" :: "r"(addr) : "memory");
}

// Allocates a 512-entry 64bit page table/directory/whatever (zeroed)
static uint64_t* alloc_page_table()
{
    uint64_t* virt = (uint64_t*)PHYS_TO_VIRT(pmm_alloc());

    for (size_t i=0; i<512; i++)
    {
        virt[i] = 0;
    }
    return virt;
}

// Kernel paging root table, that will be placed in cr3
__attribute__((aligned(4096)))
uint64_t *kernel_pml4;

// Map a page, taking virt and phys address. This will go through the paging structures
// beginning at the given root table, translate the virtual address in indexes in
// page table/directories, and then mapping the correct page table entry with the
// given physical address + flags
void paging_map_page(uint64_t* root_table, uint64_t virt, uint64_t phys, uint64_t flags)
{
    virt = PAGE_ALIGN_DOWN(virt);
    phys = PAGE_ALIGN_DOWN(phys);

    // Translate the virt address into page table indexes
    uint64_t pml4_i = PML4_INDEX(virt);
    uint64_t pdpt_i = PDPT_INDEX(virt);
    uint64_t pd_i = PD_INDEX(virt);
    uint64_t pt_i = PT_INDEX(virt);

    uint64_t *pdpt, *pd, *pt;

    // PML4
    // If the entry at index is not present, allocate enough space for it
    // then populate the entry with correct addr + flags
    if (!(root_table[pml4_i] & PTE_PRESENT))
    {
        pdpt = alloc_page_table();
        root_table[pml4_i] = VIRT_TO_PHYS(pdpt) | PTE_PRESENT | PTE_WRITABLE;
    }
    else {
        pdpt = (uint64_t *)PHYS_TO_VIRT(root_table[pml4_i] & ~0xFFFULL);
    }

    // PDPT: same here
    if (!(pdpt[pdpt_i] & PTE_PRESENT))
    {
        pd = alloc_page_table();
        pdpt[pdpt_i] = VIRT_TO_PHYS(pd) | PTE_PRESENT | PTE_WRITABLE;
    } 
    else {
        pd = (uint64_t *)PHYS_TO_VIRT(pdpt[pdpt_i] & ~0xFFFULL);
    }

    // PD: and here
    if (!(pd[pd_i] & PTE_PRESENT))
    {
        pt = alloc_page_table();
        pd[pd_i] = VIRT_TO_PHYS(pt) | PTE_PRESENT | PTE_WRITABLE;
    }
    else {
        pt = (uint64_t *)PHYS_TO_VIRT(pd[pd_i] & ~0xFFFULL);
    }

    // PT: finally, populate the page table entry
    pt[pt_i] = phys | flags | PTE_PRESENT;

    // Flush TLB (apply changes)
    invlpg((void *)virt);
}

uint64_t kernel_phys_base;
uint64_t kernel_virt_base;

void paging_init(struct limine_kernel_address_response* kaddr, struct limine_framebuffer* fb)
{
    // We should map the kernel, GDT, IDT, stack, framebuffer.
    // Optionally we could map ACPI tables (we can find them in the Limine memmap)

    kernel_phys_base = kaddr->physical_base;
    kernel_virt_base = kaddr->virtual_base;

    DEBUG("Kernel lives at virt=0x%p phys=0x%p", kernel_virt_base, kernel_phys_base);

    kernel_pml4 = alloc_page_table();

    // for debug
    uint64_t page_count = 0;

    // HHDM map first 1 GB using given offset
    for (uint64_t i=0; i<0x40000000; i += PAGE_SIZE)
    {
        //paging_kmap_page(i+hhdm_off, i, PTE_WRITABLE);
        paging_map_page(kernel_pml4, i+hhdm_off, i, PTE_WRITABLE);
        page_count++;
    }
    DEBUG("Mapped %u pages for first 1GB (HHDM)", page_count); page_count = 0;

    // Map the kernel (according to virt/phys_base given by Limine)
    // SOME DAY when we want a safer kernel we should map .text as Read/Exec
    // .rodata as Read and .data as Read/Write
    // For now who gives a shit, let's RWX all kernel
    for (uint64_t i = 0; i < KERNEL_SIZE; i += PAGE_SIZE)
    {
        //paging_kmap_page(kernel_virt_base+i, kernel_phys_base+i, PTE_WRITABLE);
        paging_map_page(kernel_pml4, kernel_virt_base+i, kernel_phys_base+i, PTE_WRITABLE);
        page_count++;
    }
    DEBUG("Mapped %u pages for kernel", page_count); page_count = 0;

    // Get the framebuffer phys/virt address, and size
    uint64_t fb_virt = (uint64_t)fb->address;
    uint64_t fb_phys = VIRT_TO_PHYS(fb_virt);
    uint64_t fb_size = fb->pitch * fb->height;
    uint64_t fb_pages = (fb_size + PAGE_SIZE-1)/PAGE_SIZE;

    // Map the framebuffer (with cache-disable & write-through)
    for (uint64_t i=0; i<fb_pages; i++)
    {
        //paging_kmap_page(fb_virt+i*PAGE_SIZE, fb_phys+i*PAGE_SIZE, PTE_WRITABLE | PTE_PCD | PTE_PWT);
        paging_map_page(kernel_pml4, fb_virt+i*PAGE_SIZE, fb_phys+i*PAGE_SIZE, PTE_WRITABLE | PTE_PCD | PTE_PWT);
        page_count++;
    }
    DEBUG("Mapped %u pages for framebuffer", page_count);

    // Finally, we load the physical address of our PML4 (root table) into cr3
    load_cr3(VIRT_TO_PHYS(kernel_pml4));
    DEBUG("cr3 loaded, we're still alive");
}