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

#include <hal/mmu.h>
#include <hal/page.h>
#include <machine/vas.h>
#include <mm/pmm.h>
#include <mm/vmm.h>
#include <stdef.h>

/*
 * Page-Table Entry (PTE) flags
 *
 * See Intel SDM Vol 3A, Section 4.5, Table 4-19
 */
#define PTE_ADDR_MASK   0x000FFFFFFFFFF000
#define PTE_P           BIT(0)        /* Present */
#define PTE_RW          BIT(1)        /* Writable */
#define PTE_US          BIT(2)        /* User r/w allowed */
#define PTE_PWT         BIT(3)        /* Page-level write-through */
#define PTE_PCD         BIT(4)        /* Page-level cache disable */
#define PTE_ACC         BIT(5)        /* Accessed */
#define PTE_DIRTY       BIT(6)        /* Dirty (written-to page) */
#define PTE_PS          BIT(7)        /* Page size */
#define PTE_GLOBAL      BIT(8)        /* Global; sticky */
#define PTE_NX          BIT(63)       /* Execute-disable */

/* 57-bit linear addresses */
#define CR4_LA57 BIT(12)

/*
 * Valid page levels
 */
typedef enum {
    PAGE_LEVEL_PML1,
    PAGE_LEVEL_PML2,
    PAGE_LEVEL_PML3,
    PAGE_LEVEL_PML4,
    PAGE_LEVEL_PML5
} PAGE_LEVEL;

/*
 * Convert protection flags to PTE flags
 */
static inline UQUAD
MmuProtToPte(UCHAR Prot)
{
   switch (Prot) {
   case PAGE_READONLY:
        return PTE_P | PTE_NX;
   case PAGE_READWRITE:
        return PTE_P | PTE_RW | PTE_NX;
   case PAGE_READWRITE_X:
        return PTE_P | PTE_RW;
   case PAGE_READONLY_X:
        return PTE_P;
   }

   return 0;
}

/*
 * Returns true if the given pagesize is valid
 *
 * @PageSize: Pagesize to check
 */
static inline BOOLEAN
MmuPageSizeValid(MMU_PAGESIZE PageSize)
{
    switch (PageSize) {
    case PAGESIZE_4K:
        return true;
    default:
        return false;
    }
}

/*
 * Obtain the toplevel index
 */
static inline PAGE_LEVEL
MmuTopLevel(VOID)
{
    UPTR Cr4;

    ASMV(
        "mov %%cr4, %0"
        : "=r" (Cr4)
        :
        : "memory"
    );

    return ISSET(Cr4, CR4_LA57)
        ? PAGE_LEVEL_PML5
        : PAGE_LEVEL_PML4;
}

/*
 * Invalidate an entry from the TLB
 *
 * @Vma: Virtual memory address to flush
 */
static inline VOID
MmuInvlpg(UPTR Vma)
{
    ASMV(
        "invlpg (%0)"
        :
        : "r" (Vma)
        : "memory"
    );
}

/*
 * Obtain the page level index
 *
 * @Vma: Virtual memory address
 * @Level: Page level
 */
static inline UQUAD
MmuLevelIndex(UPTR Vma, PAGE_LEVEL Level)
{
    switch (Level) {
    case PAGE_LEVEL_PML1: return (Vma >> 12)  & 0x1FF;
    case PAGE_LEVEL_PML2: return (Vma >> 21) & 0x1FF;
    case PAGE_LEVEL_PML3: return (Vma >> 30) & 0x1FF;
    case PAGE_LEVEL_PML4: return (Vma >> 39) & 0x1FF;
    case PAGE_LEVEL_PML5: return (Vma >> 48) & 0x1FF;
    default: return 0;
    }
}

/*
 * Extract a page table entry at a specific level
 *
 * @Vas:    Virtual address space containing top-level
 * @Vma:    Virtual memory address used as key
 * @Level:  Level to extract
 * @Alloc:  If true, allocate unset entries
 */
static UPTR *
MmuExtractLevel(MMU_VAS *Vas, UPTR Vma, PAGE_LEVEL Level, BOOLEAN Alloc)
{
    MM_PFN AllocPfn;
    UPTR Pma, *PageMap;
    PAGE_LEVEL CurrentLevel;
    USHORT Index;

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

    CurrentLevel = MmuTopLevel();
    PageMap = PMA_TO_VMA((Vas->Cr3 & PTE_ADDR_MASK));

    while (CurrentLevel > Level) {
        Index = MmuLevelIndex(Vma, CurrentLevel);

        if (ISSET(PageMap[Index], PTE_P)) {
            Pma = PageMap[Index] & PTE_ADDR_MASK;
            PageMap = PMA_TO_VMA(Pma);
            --CurrentLevel;
            continue;
        }

        if (!Alloc) {
            return NULL;
        }

        AllocPfn = MmRequestFrame();
        if (AllocPfn == 0) {
            return NULL;
        }

        Pma = AllocPfn << LOG2_PAGESIZE;
        PageMap[Index] = Pma | (PTE_P | PTE_RW | PTE_US);

        PageMap = PMA_TO_VMA(Pma);
        --CurrentLevel;
    }

    return PageMap;
}

ST_STATUS
HalMmuForkVas(MMU_VAS *Old, MMU_VAS *NewResult)
{
    UPTR *OldVma, *NewVma;
    UPTR NewPma;
    MM_PFN NewPfn;

    if (Old == NULL || NewResult == NULL) {
        return STATUS_INVALID_PARAM;
    }

    NewPfn = MmRequestFrame();
    if (NewPfn == 0) {
        return STATUS_NO_MEMORY;
    }

    NewPma = NewPfn << LOG2_PAGESIZE;
    NewVma = PMA_TO_VMA(NewPma);
    OldVma = PMA_TO_VMA((Old->Cr3 & PTE_ADDR_MASK));

    /* Start breeding the new VAS~ Fill it deep~ */
    for (USHORT Off = 0; Off < 512; ++Off) {
        if (Off < 256) {
            NewVma[Off] = 0;
        } else {
            NewVma[Off] = OldVma[Off];
        }
    }

    NewResult->Cr3 = NewPma;
    return STATUS_SUCCESS;
}

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

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

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

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

ST_STATUS
HalMmuMapSingle(MMU_VAS *Vas, UPTR Vma, UPTR Pma, USHORT Flags, MMU_PAGESIZE PagSize)
{
    UPTR *PageTable;
    USIZE PteFlags;
    USHORT Index;

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

    PageTable = MmuExtractLevel(Vas, Vma, PAGE_LEVEL_PML1, true);
    if (PageTable == NULL) {
        return STATUS_NO_MEMORY;
    }

    PteFlags = MmuProtToPte(Flags);
    Index = MmuLevelIndex(Vma, PAGE_LEVEL_PML1);

    /* Map the entry and flush the TLB */
    PageTable[Index] = Pma | PteFlags;
    MmuInvlpg(Vma);
    return STATUS_SUCCESS;
}

ST_STATUS
HalMmuUnmapSingle(MMU_VAS *Vas, UPTR Vma, MMU_PAGESIZE PageSize)
{
    UPTR *PageTable;
    USHORT Index;

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

    PageTable = MmuExtractLevel(Vas, Vma, PAGE_LEVEL_PML1, false);
    if (PageTable == NULL) {
        return STATUS_NOT_FOUND;
    }

    /* Unmap the page */
    Index = MmuLevelIndex(Vma, PAGE_LEVEL_PML1);
    PageTable[Index] = 0;
    MmuInvlpg(Vma);
    return STATUS_SUCCESS;
}