/*
 * @author  xamidev <xamidev@riseup.net>
 * @brief   x64 4-level paging implementation
 * @license GPL-3.0-only
 */

#include <mem/paging.h>
#include <mem/pmm.h>
#include <kernel.h>
#include <stddef.h>
#include <limine.h>
#include <config.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
*/

/*
 * load_cr3 - Load a new value into the CR3 register
 * @value: the value to load
 *
 * This function is used to load the physical address
 * of the root page table (PML4), to switch the paging
 * structures the CPU sees and uses.
 */
void load_cr3(uint64_t value) {
    asm volatile ("mov %0, %%cr3" :: "r"(value) : "memory");
}

/*
 * invlpg - Invalidates a Translation Lookaside Buffer entry
 * @addr: page memory address
 *
 * This function is used to flush at least the TLB entrie(s)
 * for the page that contains the <addr> address.
 */
static inline void invlpg(void *addr)
{
    asm volatile("invlpg (%0)" :: "r"(addr) : "memory");
}

/*
 * alloc_page_table - Page table allocation
 *
 * This function allocates enough memory for a 512-entry
 * 64-bit page table, for any level (PML4/3/2).
 *
 * Memory allocated here is zeroed.
 *
 * Return:
 * <virt> - Pointer to allocated page table
 */
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;

/*
 * paging_map_page - Mapping a memory page
 * @root_table: Address of the PML4
 * @virt:       Virtual address
 * @phys:       Physical address
 * @flags:      Flags to set on page
 *
 * This function maps the physical address <phys> to the virtual
 * address <virt>, using the paging structures beginning at
 * <root_table>. <flags> can be set according to the PTE_FLAGS enum.
 *
 * If a page table/directory entry is not present yet, it creates it.
 */
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] & PTE_ADDR_MASK);
    }

    // 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] & PTE_ADDR_MASK);
    }

    // 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] & PTE_ADDR_MASK);
    }

    // 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;

/*
 * paging_init - Paging initialization
 * @boot_ctx: Boot context structure
 *
 * This function initializes new paging structures, to replace
 * the ones given by the bootloader.
 *
 * It maps the kernel, the HHDM space, and the framebuffer.
 */
void paging_init(struct boot_context boot_ctx)
{
    // 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 = boot_ctx.kaddr->physical_base;
    kernel_virt_base = boot_ctx.kaddr->virtual_base;
    struct limine_framebuffer* fb = boot_ctx.fb;

    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;

    // Find max physical address from limine memmap
    uint64_t max_phys = 0;
    for (uint64_t i=0; i<boot_ctx.mmap->entry_count; i++) {
        struct limine_memmap_entry* entry = boot_ctx.mmap->entries[i];
        if (entry->length == 0) {
            continue;
        }
        uint64_t top = entry->base + entry->length;
        if (top > max_phys) {
            max_phys = top;
        }
    }

    // 4GB
    if (max_phys > PAGING_MAX_PHYS) {
        DEBUG("WARNING: max_phys capped to 4GB (%x) (from max_phys=%p)", PAGING_MAX_PHYS, max_phys);
        max_phys = PAGING_MAX_PHYS;
    }

    // HHDM map up to max_phys or PAGING_MAX_PHYS, whichever is smaller, using given offset
    for (uint64_t i=0; i<max_phys; i += PAGE_SIZE) {
        paging_map_page(kernel_pml4, i+hhdm_off, i, PTE_WRITABLE | PTE_PRESENT);
        page_count++;
    }
    DEBUG("Mapped %u pages up to 0x%p (HHDM)", page_count, max_phys); 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_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_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("Loaded kernel PML4 into CR3");
}