/*
 * Copyright (c) 2026, Chloe M.
 * Provided under the BSD-3 clause.
 *
 * Description: Platform MMU HAL interface
 * Author:      Chloe M.
 */

#include <hal/mmu.h>
#include <machine/mmu.h>
#include <machine/tlb.h>
#include <mm/pm.h>
#include <mm/vm.h>
#include <ke/bugcheck.h>
#include <ptdef.h>

/*
 * Valid paging structure levels
 */
typedef enum {
    PAGEMAP_PML1,
    PAGEMAP_PML2,
    PAGEMAP_PML3,
    PAGEMAP_PML4
} PAGEMAP_LEVEL;

/*
 * Returns true if the given page size is valid
 *
 * @PageSize: Page size to check
 */
ALWAYS_INLINE static inline BOOLEAN
MmuPageSizeValid(MMU_PAGESIZE PageSize)
{
    switch (PageSize) {
    case PAGESIZE_4K:
        return true;
    }

    return false;
}

/*
 * Obtain the pagemap top-level index
 *
 * XXX: In future revisions we should read CR4 and determine
 *      if 57-bit linear addresses are enabled for 5-level
 *      paging.
 */
ALWAYS_INLINE static inline PAGEMAP_LEVEL
MmuTopLevel(VOID)
{
    return PAGEMAP_PML4;
}

/*
 * Obtain PTE flags from MI page protection
 * flags
 *
 * @Prot:  Protection flags to convert
 */
static USIZE
MmuProtToPte(MM_PROT Prot)
{
    USIZE PteFlags = PTE_P | PTE_NX;

    if (ISSET(Prot, PAGE_WRITE))
        PteFlags |= PTE_RW;
    if (ISSET(Prot, PAGE_EXEC))
        PteFlags &= ~PTE_NX;
    if (ISSET(Prot, PAGE_USER))
        PteFlags |= PTE_US;

    return PteFlags;
}

/*
 * Extract a pagemap index from a virtual address for a given
 * level.
 *
 * @VirtualBase:  Virtual base address
 * @Level:        Pagemap level to get index of
 *
 * XXX: Caller should ensure the page alignment of @VirtualBase
 */
static inline USHORT
MmuPageMapIndex(UPTR VirtualBase, PAGEMAP_LEVEL Level)
{
    switch (Level) {
    case PAGEMAP_PML1:
        return (VirtualBase >> 12) & 0x1FF;
    case PAGEMAP_PML2:
        return (VirtualBase >> 21) & 0x1FF;
    case PAGEMAP_PML3:
        return (VirtualBase >> 30) & 0x1FF;
    case PAGEMAP_PML4:
        return (VirtualBase >> 39) & 0x1FF;
    }

    /* This should not happen */
    KeBugCheck(
        BUGCHECK_IO_ERROR,
        "bad pagemap level in MmuPageMapIndex()\n"
    );
}

/*
 * Extract a pagemap at a specific level
 *
 * @Vas:            Virtual address space to translate from
 * @VirtualBase:    Virtual base address
 * @Level:          Level to extract
 * @Alloc:          If true, allocate new level
 *
 * Returns the pogemap level base on success, otherwise NULL
 * on failure.
 *
 * XXX: Caller should ensure the page alignment of @VirtualBase
 */
static UPTR *
MmuPageMapExtract(MMU_VAS *Vas, UPTR VirtualBase, PAGEMAP_LEVEL Level, BOOLEAN Alloc)
{
    UPTR *PageMap;
    PAGEMAP_LEVEL CurrentLevel;
    MM_PFN Pfn;
    UPTR PhysicalBase;
    USIZE Index;

    if (Vas == NULL) {
        return NULL;
    }

    /* Get the top-level */
    CurrentLevel = MmuTopLevel();
    PageMap = PMA_TO_VMA((Vas->Cr3 & PTE_ADDR_MASK));

    /* Walk through each level */
    while (CurrentLevel > Level) {
        Index = MmuPageMapIndex(VirtualBase, CurrentLevel);

        /* Is this entry already present? */
        if (ISSET(PageMap[Index], PTE_P)) {
            PageMap = PMA_TO_VMA((PageMap[Index] & PTE_ADDR_MASK));
            --CurrentLevel;
            continue;
        }

        if (!Alloc) {
            return NULL;
        }

        /*
         * Now we allocate a level for this entry
         *
         * XXX: We need to consider how we'll clean up when unmapping pages
         *      as the destruction of a virtual address space will end up
         *      leaking memory if this point is reached.
         */
        Pfn = MmRequestFrame();
        if (Pfn == PFN_ERROR) {
            return NULL;
        }

        PhysicalBase = PFN_TO_ADDRESS(Pfn);
        PageMap[Index] = PhysicalBase | (PTE_P | PTE_RW | PTE_US);
        PageMap = PMA_TO_VMA(PhysicalBase);
        --CurrentLevel;
    }

    return PageMap;
}

VOID
HalMmuReadVas(MMU_VAS *Result)
{
    if (Result == NULL) {
        return;
    }

    ASMV(
        "mov %%cr3, %0"
        : "=r" (Result->Cr3)
        :
        : "memory"
    );
}

VOID
HalMmuWriteVas(const MMU_VAS *Vas)
{
    if (Vas == NULL) {
        return;
    }

    ASMV(
        "mov %0, %%cr3"
        :
        : "r" (Vas->Cr3)
        : "memory"
    );
}

PT_STATUS
HalMmuMapPage(
    MMU_VAS *Vas, UPTR VirtualBase,
    UPTR PhysicalBase, MM_PROT Prot,
    MMU_PAGESIZE PageSize)
{
    UPTR *PageTable;
    USIZE PageTableIndex;
    USIZE PteFlags;

    if (Vas == NULL) {
        return STATUS_INVALID_PARAM;
    }

    if (!MmuPageSizeValid(PageSize)) {
        return STATUS_INVALID_PARAM;
    }

    PageTable = MmuPageMapExtract(Vas, VirtualBase, PAGEMAP_PML1, true);
    if (PageTable == NULL) {
        return STATUS_NO_MEMORY;
    }

    PageTableIndex = MmuPageMapIndex(VirtualBase, PAGEMAP_PML1);
    PteFlags = MmuProtToPte(Prot);

    /* Map the memory and flush the TLB entry */
    PageTable[PageTableIndex] = PhysicalBase | PteFlags;
    MdTlbFlushSingle(VirtualBase);
    return STATUS_SUCCESS;
}