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

#include <mm/vmm.h>
#include <mm/vad.h>
#include <mm/pmm.h>
#include <hal/mmu.h>
#include <hal/page.h>
#include <ex/trace.h>
#include <ke/knot.h>
#include <stapi/status.h>
#include <string.h>

#define VALLOC_MAX_PAGES 8192
#define VALLOC_MAX (VALLOC_BASE + (VALLOC_MAX_PAGES * PAGESIZE))
#define PAGES_PER_GRAB 16

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

static MMU_VAS KernelVas;
static MM_VAD_LIST VadList;
static UPTR BumpPtr;

/*
 * TODO: Add some locking around this
 */
static ST_STATUS
MmGrabPages(VOID)
{
    USIZE MaxOff;
    UPTR AllocBase;
    MMU_VAS Vas;
    MM_PFN AllocPfn;
    ST_STATUS Status;

    MaxOff = PAGES_PER_GRAB * PAGESIZE;
    HalMmuReadVas(&Vas);

    for (USIZE Off = 0; Off < MaxOff; Off += PAGESIZE) {
        /* Don't blast past the limit */
        if ((BumpPtr + Off) >= VALLOC_MAX) {
            return STATUS_NO_MEMORY;
        }

        AllocBase = BumpPtr + Off;
        AllocPfn = MmRequestFrame();
        if (AllocPfn == 0) {
            return STATUS_NO_MEMORY;
        }

        Status = HalMmuMapSingle(
            &Vas,
            AllocBase,
            AllocPfn << LOG2_PAGESIZE,
            PAGE_READWRITE,
            PAGESIZE_4K
        );

        if (Status != STATUS_SUCCESS) {
            return Status;
        }

        MmVadListAppend(&VadList, VAD_FROM_VMA(AllocBase));
    }

    BumpPtr = AllocBase;
    return STATUS_SUCCESS;
}

/*
 * TODO: Consider locking
 */
VOID *
MmAllocPages(USIZE Count)
{
    MM_VAD *First, *Vad;
    ST_STATUS Status;

    if (Count == 0) {
        return NULL;
    }

    First = NULL;

    /*
     * Request a VAD from the allocation VAD list and pull in
     * more pages if we need to.
     *
     * XXX: We need to consider how we'll deal with fragmentation
     *      here, once it fills up it may become fragmented.
     */
    for (USIZE Iter = 0; Iter < Count; ++Iter) {
        Vad = MmVadListPop(&VadList);
        if (Vad == NULL) {
            Status = MmGrabPages();
            if (Status != STATUS_SUCCESS)
                return NULL;

            Vad = MmVadListPop(&VadList);
        }

        /* Still no pages? */
        if (Vad == NULL) {
            return NULL;
        }

        if (First == NULL) {
            First = Vad;
        }
    }

    RtlMemSet(First, 0, sizeof(*First));
    return (VOID *)First;
}

VOID
MmInitVmm(VOID)
{
    MMU_VAS OldVas;
    ST_STATUS Status;

    DTRACE("breeding kvas from bootstrap instance...\n");
    HalMmuReadVas(&OldVas);
    Status = HalMmuForkVas(&OldVas, &KernelVas);

    if (Status != STATUS_SUCCESS) {
        KeKnot(KNOT_OOM, "could not fork bootstrap vas\n");
    }

    HalMmuWriteVas(&KernelVas);
    DTRACE("ah!~... [ok @ %p]\n", VAS_BASE(&KernelVas));
    BumpPtr = VALLOC_BASE;
}