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

#include <mm/vm.h>
#include <mm/pm.h>
#include <mm/vad.h>
#include <ex/trace.h>
#include <hal/page.h>
#include <hal/mmu.h>
#include <ke/bugcheck.h>
#include <ptdef.h>

#define DTRACE(Fmt, ...) \
    TRACE("[ VM ]: " Fmt, ##__VA_ARGS__)

/*
 * Number of pages we start with that can be valloc'd
 * total.
 */
#define INITIAL_VALLOC_PAGES 128

/* Use to safely convert page size */
#define PAGE_SIZE(PS)                \
    ((PS) >= NELEM(PageSizeTable))   \
        ? PageSizeTable[PAGESIZE_4K] \
        : PageSizeTable[(PS)]

/* Page size table */
static USIZE PageSizeTable[] = {
    [PAGESIZE_4K] = 0x1000
};

/* Kernel VAS */
static MMU_VAS KVas;

/* Bump pointer and vad tree */
static UPTR VallocBump = VALLOC_BASE;
static MM_VAD_TREE VadTree;

PT_STATUS
MmAllocatePages(USIZE Count, MM_RANGE *Result)
{
    MM_VAD *Vad;

    if (Count == 0 || Result == NULL) {
        return STATUS_INVALID_PARAM;
    }

    Vad = MmVadTreeAllocate(&VadTree, Count);
    if (Vad == NULL) {
        return STATUS_NO_MEMORY;
    }

    *Result = Vad->Range;
    return STATUS_SUCCESS;
}

PT_STATUS MmUnmapRange(MMU_VAS *Vas, const MM_RANGE *Range, MMU_PAGESIZE PageSize)
{
    PT_STATUS Status;
    UPTR Off = 0;
    UPTR VirtualBase;
    USIZE Gran;

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

    /* Align down the base addresses */
    VirtualBase = ALIGN_DOWN(Range->VirtualBase, PAGESIZE);
    Gran = PAGE_SIZE(PageSize);

    for (Off = 0; Off < Range->Length; Off += Gran) {
        Status = HalMmuUnmapPage(
            Vas,
            VirtualBase + Off,
            PageSize
        );

        if (PT_ERROR(Status)) {
            return Status;
        }
    }

    return STATUS_SUCCESS;
}

PT_STATUS
MmMapRange(MMU_VAS *Vas, const MM_RANGE *Range, MM_PROT Prot, MMU_PAGESIZE PageSize)
{
    PT_STATUS Status;
    UPTR Off = 0;
    UPTR VirtualBase;
    UPTR PhysicalBase;
    USIZE Gran;

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

    /* Align down the base addresses */
    VirtualBase = ALIGN_DOWN(Range->VirtualBase, PAGESIZE);
    PhysicalBase = ALIGN_DOWN(Range->PhysicalBase, PAGESIZE);
    Gran = PAGE_SIZE(PageSize);

    for (Off = 0; Off < Range->Length; Off += Gran) {
        Status = HalMmuMapPage(
            Vas,
            VirtualBase + Off,
            PhysicalBase + Off,
            Prot,
            PageSize
        );

        if (PT_ERROR(Status)) {
            MmUnmapRange(Vas, Range, PageSize);
            return Status;
        }
    }

    return STATUS_SUCCESS;
}

VOID
MmVmInit(VOID)
{
    PT_STATUS Status;
    MMU_VAS VasCurrent;
    UPTR VirtualStart;
    UPTR VirtualEnd, Off;
    UPTR PhysicalBase;
    USIZE NumPages;
    MM_PFN Pfn;
    MM_RANGE *CurrentRange;
    MM_VAD *CurrentVad;

    HalMmuReadVas(&VasCurrent);
    MmVadTreeInit(&VadTree);
    Status = HalMmuForkVas(&VasCurrent, &KVas);
    if (PT_ERROR(Status)) {
        KeBugCheck(
            BUGCHECK_OOM,
            "failed to allocate kernel vas\n"
        );
    }

    HalMmuWriteVas(&KVas);
    VirtualStart = VallocBump;
    VallocBump += INITIAL_VALLOC_PAGES * PAGESIZE;
    VirtualEnd = VallocBump;
    NumPages = (VirtualEnd - VirtualStart) / PAGESIZE;

    for (Off = 0; Off < NumPages * PAGESIZE; Off += PAGESIZE) {
        Pfn = MmRequestFrame();
        if (Pfn == 0) {
            KeBugCheck(BUGCHECK_OOM, "out of memory\n");
        }

        PhysicalBase = PFN_TO_ADDRESS(Pfn);
        Status = HalMmuMapPage(
            &KVas,
            VirtualStart + Off,
            PhysicalBase,
            PAGE_READ| PAGE_WRITE,
            PAGESIZE_4K
        );

        if (PT_ERROR(Status)) {
            KeBugCheck(BUGCHECK_OOM, "failed to map valloc pages\n");
        }

        CurrentVad = (MM_VAD *)(VirtualStart + Off);
        CurrentRange = &CurrentVad->Range;
        CurrentRange->VirtualBase = VirtualStart + Off;
        CurrentRange->PhysicalBase = PhysicalBase;
        CurrentRange->Length = PAGESIZE;
        Status = MmVadTreeInsert(
            &VadTree,
            CurrentVad
        );

        if (PT_ERROR(Status)) {
            KeBugCheck(BUGCHECK_MISC, "failed to initialize vad tree\n");
        }
    }

    DTRACE("initialized %d valloc pages\n", NumPages);
}