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

#include <ex/trace.h>
#include <ke/bpal.h>
#include <mm/pmm.h>
#include <mm/vmm.h>
#include <hal/page.h>
#include <units.h>
#include <stdef.h>

#define LOAD_DELAY 128

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

/* Use to convert type constants to strings */
#define MEM_TYPE(TYPE)          \
    ((TYPE) < NELEM(TypeTab))   \
        ? TypeTab[(TYPE)]       \
        : "invalid"

/*
 * Physical frame descriptor
 *
 * @Next:  Next PFD in list
 * @Prev:  Previous PFD in list
 */
typedef struct _MM_PFD {
    struct _MM_PFD *Next;
    struct _MM_PFD *Prev;
} MM_PFD;

/*
 * Page frame descriptor list
 *
 * @First:  First entry in list
 * @Last:   Last entry in list
 */
typedef struct {
    MM_PFD *First;
    MM_PFD *Last;
} MM_PFD_LIST;

/*
 * Memory type constant to human readable
 * string table.
 */
static const char *TypeTab[] = {
    [MEMORY_USABLE]         = "usable",
    [MEMORY_RESERVED]       = "reserved",
    [MEMORY_ACPI_RECLAIM]   = "acpi reclaimable",
    [MEMORY_ACPI_NVS]       = "acpi nvs",
    [MEMORY_BAD]            = "bad*",
    [MEMORY_BOOTLOADER]     = "bootloader",
    [MEMORY_KERNEL]         = "stoskrnl.sys",
    [MEMORY_FRAMEBUFFER]    = "framebuffer",
    [MEMORY_ACPI_TABLES]    = "acpi tables"
};

/* Globals */
static USIZE UsableMemory = 0;
static MM_PFD_LIST PfdList;

/*
 * Append a page frame descriptor to a page frame descriptor list
 *
 * @List: PFD list to append to
 * @Pfd:  PFD to append
 */
static VOID
MmPfdAppend(MM_PFD_LIST *List, MM_PFD *Pfd)
{
    MM_PFD *Last;

    if (List == NULL || Pfd == NULL) {
        return;
    }

    if (List->Last == NULL || List->First == NULL) {
        List->Last = Pfd;
        List->First = Pfd;
    } else {
        Last = List->Last;
        Pfd->Prev = Last;
        Last->Next = Pfd;
        List->Last = Pfd;
    }
}

/*
 * Pop a page frame descriptor from a page frame descriptor list
 *
 * @List: Page frame descriptor list to pop from
 *
 * Returns NULL if there are no entries left
 */
static MM_PFD *
MmPfdPop(MM_PFD_LIST *List)
{
    MM_PFD *Last, *Prev;

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

    Last = List->Last;
    if (Last == NULL) {
        return NULL;
    }

    /* Unlink from the list */
    Prev = Last->Prev;
    List->Last = List->Last->Prev;
    Prev->Next = NULL;
    List->Last = Prev;

    return Last;
}

/*
 * Print the units of a size in a pretty format
 *
 * @Length: Number of bytes
 * @Title:  Title of units
 */
static inline VOID
MmPrintUnits(USIZE Length, const CHAR *Title)
{
    if (Title == NULL) {
        return;
    }

    if (Length >= UNIT_GIB) {
        DTRACE("%d gib %s\n", Length / UNIT_GIB, Title);
    } else if (Length >= UNIT_MIB) {
        DTRACE("%d mib %s\n", Length / UNIT_MIB, Title);
    } else {
        DTRACE("%010d b %s\n", Length, Title);
    }
}

/*
 * Display a loading status
 */
static VOID
MmLoadUpdate(VOID)
{
    const UCHAR LoadMax = 3;
    static UCHAR LoadIdx = 0;
    static CHAR LoadProg[3] = {'/', '-', '\\'};

    TRACE("\x1b[?25l%c \b\b", LoadProg[LoadIdx++]);
    LoadIdx %= LoadMax;
}

static VOID
MmProbeMemory(VOID)
{
    BPAL_HANDLE BpalHandle;
    MEMMAP_ENTRY Entry;
    UPTR EntryEnd;
    VOID *Vma, *Base;
    USIZE Idx, Offset;
    USIZE EntryCount = 0;
    ST_STATUS Status;

    KeBpalGetHandle(&BpalHandle);
    for (Idx = 0;; ++Idx) {
        Status = BpalHandle.MemEntryIdx(Idx, &Entry);
        if (Status != STATUS_SUCCESS) {
            break;
        }

        EntryEnd = Entry.Base + Entry.Length;
        DTRACE("[%p - %p) ~ %s\n", Entry.Base, EntryEnd, MEM_TYPE(Entry.Type));

        /* Only usable entries from this point on */
        if (Entry.Type != MEMORY_USABLE) {
            continue;
        }

        Base = PMA_TO_VMA(Entry.Base);
        for (Offset = 0; Offset < Entry.Length; Offset += PAGESIZE) {
            Vma = PTR_OFFSET(Base, Offset);
            MmPfdAppend(&PfdList, Vma);
            EntryCount++;

            /* Display a little loader as this takes awhile */
            if ((EntryCount & (LOAD_DELAY - 1)) == 0) {
                MmLoadUpdate();
            }
        }

        UsableMemory += Entry.Length;
    }

    TRACE("\b\x1b[?25h");
    DTRACE("%d entries populated\n", EntryCount);
    MmPrintUnits(UsableMemory, "avl");
}

VOID
MmInitPmm(VOID)
{
    DTRACE("sniffing out installed ram...\n");
    MmProbeMemory();
}