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

#include <mm/pm.h>
#include <mm/vm.h>
#include <ke/bpal.h>
#include <ex/trace.h>
#include <string.h>
#include <units.h>

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

/* Safe macro to get region string */
#define PM_REGION(TYPE)                 \
    ((TYPE) >= NELEM(PmRegionTable))    \
        ? "bad type"                    \
        : PmRegionTable[(TYPE)]

/*
 * Table of known memory regions
 */
static const CHAR *PmRegionTable[] = {
    [MEMORY_USABLE]         = "available",
    [MEMORY_RESERVED]       = "reserved",
    [MEMORY_ACPI_RECLAIM]   = "acpi reclaimable",
    [MEMORY_ACPI_NVS]       = "acpi nv-sleep",
    [MEMORY_BAD]            = "bad",
    [MEMORY_BOOTLOADER]     = "bootloader",
    [MEMORY_KERNEL]         = "system",
    [MEMORY_FRAMEBUFFER]    = "vram",
    [MEMORY_ACPI_TABLES]    = "acpi tables"
};

/*
 * A page frame descriptor is an entry within the physical memory
 * freelist.
 *
 * @Next:  Next PFD in freelist
 * @Prev:  Previous PFD in freelist
 */
typedef struct _MM_PFD {
    struct _MM_PFD *Next;
    struct _MM_PFD *Prev;
} MM_PFD;

/*
 * A page frame descriptor list holds many page frame
 * descriptors and marks the start and end of a physical
 * memory freelist.
 *
 * @First: First entry in freelist
 * @Last:  Last entry in freelist
 *
 * TODO: Add a lock to this structure
 */
typedef struct {
    MM_PFD *First;
    MM_PFD *Last;
} MM_PFD_LIST;

/* Standard usable memory list */
static MM_PFD_LIST AvlList;

/* Memory statistics */
static USIZE UsableMemory;

/*
 * Used to convert bytes into human readable units
 */
static inline VOID
PmZoneSize(const CHAR *TypeName, USIZE Length)
{
    if (TypeName == NULL) {
        return;
    }

    if (Length >= UNIT_GIB) {
        DTRACE("%s : %d GiB\n", TypeName, Length / UNIT_GIB);
    } else if (Length >= UNIT_MIB) {
        DTRACE("%s : %d MiB\n", TypeName, Length / UNIT_MIB);
    } else if (Length >= UNIT_KIB) {
        DTRACE("%s : %d KiB\n", TypeName, Length / UNIT_KIB);
    } else {
        DTRACE("%s : %d bytes\n", TypeName, Length);
    }
}

/*
 * Appends a page frame descriptor to a page frame descriptor
 * list.
 *
 * @List: Page frame descriptor list to append to
 * @Pfd:  Page frame descriptor to append
 */
static VOID
PmPfdListAppend(MM_PFD_LIST *List, MM_PFD *Pfd)
{
    MM_PFD *LastPfd;

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

    Pfd->Next = NULL;
    Pfd->Prev = NULL;

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

/*
 * Pop a page frame descriptor from the beginning of a page
 * frame descriptor list.
 *
 * @List: Page frame descriptor list to pop from
 *
 * Returns NULL on failure
 */
static MM_PFD *
PmPfdListPop(MM_PFD_LIST *List)
{
    MM_PFD *FirstPfd;

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

    FirstPfd = List->First;
    List->First = FirstPfd->Next;
    return FirstPfd;
}

/*
 * Probe physical memory for usable regions and list
 * installed memory.
 */
static VOID
PmProbe(VOID)
{
    PT_STATUS Status;
    BPAL_HANDLE BpalHandle;
    MEMMAP_ENTRY Entry;
    UPTR Start, End, Off;
    MM_PFD *PfdBase, *PfdCurrent;
    USIZE Index, PfdCount = 0;
    USIZE ZoneCount = 0;

    /*
     * We obtain memory map entries from the BPAL layer as it is provided
     * by the bootloader which got it from firmware.
     */
    KeBpalGetHandle(&BpalHandle);

    /*
     * Iterate each memory entry, we use two terms here:
     *
     * Zones: A zone is a contigious region of frames and consists
     *        of many page frame descriptors.
     *
     * PFDs:  A page frame descriptor represents a single frame within
     *        a zone.
     */
    for (Index = 0;; ++Index) {
        Status = BpalHandle.MemEntryIdx(Index, &Entry);

        /* Are we out of entries to fetch? */
        if (PT_ERROR(Status)) {
            break;
        }

        Start = Entry.Base;
        End = Start + Entry.Length;
        DTRACE(
            "%p - %p : %s\n",
            Start,
            End,
            PM_REGION(Entry.Type)
        );

        /* Usable entries from this point only */
        if (Entry.Type != MEMORY_USABLE) {
            continue;
        }

        /* Append this zone */
        PfdBase = (MM_PFD *)PMA_TO_VMA(Entry.Base);
        for (Off = 0; Off < Entry.Length; Off += PAGESIZE) {
            PfdCurrent = PTR_OFFSET(PfdBase, Off);
            PmPfdListAppend(&AvlList, PfdCurrent);
            ++PfdCount;
        }

        ++ZoneCount;
        UsableMemory += Entry.Length;
    }

    /* Log some statistics */
    DTRACE("%d pfds, %d zones\n", PfdCount, ZoneCount);
    PmZoneSize("usable", UsableMemory);
}

MM_PFN
MmRequestFrame(VOID)
{
    MM_PFD *Pfd;
    UPTR PhysicalBase;

    Pfd = PmPfdListPop(&AvlList);
    RtlMemSet(Pfd, 0, PAGESIZE);

    PhysicalBase = VMA_TO_PMA(Pfd);
    return ADDRESS_TO_PFN(PhysicalBase);
}

VOID
MmReclaimFrame(MM_PFN Pfn)
{
    MM_PFD *Pfd;
    UPTR PhysicalBase;

    if (Pfn == PFN_ERROR) {
        return;
    }

    PhysicalBase = PFN_TO_ADDRESS(Pfn);
    Pfd = (VOID *)PMA_TO_VMA(PhysicalBase);
    PmPfdListAppend(&AvlList, Pfd);
}

VOID
MmPmInit(VOID)
{
    DTRACE("probing physical memory...\n");
    PmProbe();
}