// mm_phys.c - physical memory manager
//
// The free-frame stack is a flat array of PHYSICAL_ADDRESS values carved out of
// the first suitable usable region before any allocation is served.  All access
// to physical memory goes through the direct physical window established by the
// bootloader at PhysicalWindowBase (always 0xFFFF880000000000).
//
// Lock: a single spinlock guards Push/Pop. O(1) for both operations.
#include "mm_phys.h"
#include "ke_spinlock.h"
#include <stdint.h>
#include <stddef.h>

// ---------------------------------------------------------------------------
// Physical-to-virtual translation via the direct map window.
// ---------------------------------------------------------------------------
static uint64_t gPhysBase = 0;   // set by MmInitializePhysical

static inline void* PhysToVirt(PHYSICAL_ADDRESS phys) {
    return (void *)(phys + gPhysBase);
}

// ---------------------------------------------------------------------------
// Frame stack
// ---------------------------------------------------------------------------

typedef struct {
    PHYSICAL_ADDRESS* Stack;    // pointer into physical window
    size_t            Top;      // index of next free slot (stack grows up)
    size_t            Capacity;
    size_t            Total;    // total frames ever pushed
    SPINLOCK          Lock;
} FRAME_STACK;

static FRAME_STACK gFrameStack;

// Push a physical frame address onto the stack.
static void PushFrame(PHYSICAL_ADDRESS frame) {
    KeAcquireSpinlock(&gFrameStack.Lock);
    if (gFrameStack.Top < gFrameStack.Capacity) {
        gFrameStack.Stack[gFrameStack.Top++] = frame;
    }
    KeReleaseSpinlock(&gFrameStack.Lock);
}

// ---------------------------------------------------------------------------
// Initialisation
// ---------------------------------------------------------------------------
void MmInitializePhysical(const MEMORY_MAP* map, uint64_t physicalWindowBase) {
    gPhysBase = physicalWindowBase;

    // --- Phase 0: Count total usable frames so we can size the stack
    //              dynamically instead of capping at a hardcoded 1 GB. ---
    uint64_t totalUsableBytes = 0;
    for (uint32_t i = 0; i < map->RegionCount; i++) {
        if (map->Regions[i].Type == MemoryUsable)
            totalUsableBytes += map->Regions[i].Length;
    }
    const size_t maxFrames  = (size_t)(totalUsableBytes / PAGE_SIZE);
    const size_t stackBytes = maxFrames * sizeof(PHYSICAL_ADDRESS);

    // --- Phase 1: Find a usable region large enough to hold the stack. ---
    PHYSICAL_ADDRESS stackPhys = 0;
    for (uint32_t i = 0; i < map->RegionCount; i++) {
        const MEMORY_REGION* r = &map->Regions[i];
        if (r->Type == MemoryUsable && r->Length >= stackBytes + PAGE_SIZE) {
            stackPhys = r->Base;
            break;
        }
    }

    if (!stackPhys) {
        for (;;) { __halt(); }
    }

    gFrameStack.Stack    = (PHYSICAL_ADDRESS *)PhysToVirt(stackPhys);
    gFrameStack.Top      = 0;
    gFrameStack.Capacity = maxFrames;
    gFrameStack.Total    = 0;
    gFrameStack.Lock     = (SPINLOCK)SPINLOCK_INIT;

    // --- Phase 2: Push all free 4 KB frames, skipping the stack pages. ---
    const PHYSICAL_ADDRESS stackEnd =
        stackPhys + ((stackBytes + PAGE_SIZE - 1) & PAGE_MASK);

    for (uint32_t i = 0; i < map->RegionCount; i++) {
        const MEMORY_REGION* r = &map->Regions[i];
        if (r->Type != MemoryUsable) continue;

        PHYSICAL_ADDRESS base = r->Base;
        PHYSICAL_ADDRESS end  = r->Base + r->Length;

        for (PHYSICAL_ADDRESS frame = base; frame < end; frame += PAGE_SIZE) {
            if (frame >= stackPhys && frame < stackEnd) continue;
            if (gFrameStack.Top < gFrameStack.Capacity) {
                gFrameStack.Stack[gFrameStack.Top++] = frame;
                gFrameStack.Total++;
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
PHYSICAL_ADDRESS MmAllocateFrame(void) {
    KeAcquireSpinlock(&gFrameStack.Lock);
    if (gFrameStack.Top == 0) {
        KeReleaseSpinlock(&gFrameStack.Lock);
        return (PHYSICAL_ADDRESS)0;
    }
    PHYSICAL_ADDRESS frame = gFrameStack.Stack[--gFrameStack.Top];
    KeReleaseSpinlock(&gFrameStack.Lock);
    return frame;
}

void MmFreeFrame(PHYSICAL_ADDRESS frame) {
    if (!frame || (frame & (PAGE_SIZE - 1))) return;   // must be page-aligned
    PushFrame(frame);
}

PHYSICAL_ADDRESS MmAllocateContiguous(size_t count) {
    if (!count) return (PHYSICAL_ADDRESS)0;

    // Brute-force: scan the stack for 'count' consecutive frames.
    // Acceptable for small DMA buffers; not for general use.
    KeAcquireSpinlock(&gFrameStack.Lock);

    // Build a temporary sorted view - too expensive; instead do a simple
    // linear scan for a contiguous run already on the stack (lucky case),
    // or fall back to pulling frames one by one and checking adjacency.
    // For now: pull 'count' frames and verify they are contiguous.
    // If not, return them all and report failure.
    if (gFrameStack.Top < count) {
        KeReleaseSpinlock(&gFrameStack.Lock);
        return (PHYSICAL_ADDRESS)0;
    }

    // Attempt: take the top 'count' entries and check.
    PHYSICAL_ADDRESS base = gFrameStack.Stack[gFrameStack.Top - 1];
    int ok = 1;
    for (size_t i = 0; i < count; i++) {
        PHYSICAL_ADDRESS f = gFrameStack.Stack[gFrameStack.Top - 1 - i];
        if (f != base + i * PAGE_SIZE) { ok = 0; break; }
    }

    if (ok) {
        gFrameStack.Top -= count;
        KeReleaseSpinlock(&gFrameStack.Lock);
        return base;
    }

    KeReleaseSpinlock(&gFrameStack.Lock);
    return (PHYSICAL_ADDRESS)0;   // caller should retry or fail gracefully
}

size_t MmFreeFrameCount(void) {
    return gFrameStack.Top;
}

size_t MmTotalFrameCount(void) {
    return gFrameStack.Total;
}