// 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 <kernel/mm_phys.h>
#include <kernel/ke_spinlock.h>

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

static inline void* KiPhysToVirt(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 KiPushFrame(PHYSICAL_ADDRESS frame) {
    uintptr_t flags = KeAcquireSpinlock(&gFrameStack.Lock);
    if (gFrameStack.Top < gFrameStack.Capacity) {
        gFrameStack.Stack[gFrameStack.Top++] = frame;
        gFrameStack.Total++;
    }
    KeReleaseSpinlock(&gFrameStack.Lock, flags);
}

// ---------------------------------------------------------------------------
// Initialisation
// ---------------------------------------------------------------------------
bool 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) {
        return false;
    }

    gFrameStack.Stack    = (PHYSICAL_ADDRESS *)KiPhysToVirt(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;
            KiPushFrame(frame);
        }
    }

    return true;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
static PHYSICAL_ADDRESS KiPopFrameAtIndexLocked(size_t idx) {
    PHYSICAL_ADDRESS frame = gFrameStack.Stack[idx];
    size_t tail = gFrameStack.Top - idx - 1;
    for (size_t j = 0; j < tail; j++)
        gFrameStack.Stack[idx + j] = gFrameStack.Stack[idx + 1 + j];
    gFrameStack.Top--;
    return frame;
}

PHYSICAL_ADDRESS MmAllocateFrameInRange(PHYSICAL_ADDRESS minInclusive,
                                        PHYSICAL_ADDRESS maxExclusive) {
    minInclusive = (minInclusive + PAGE_SIZE - 1) & PAGE_MASK;
    maxExclusive &= PAGE_MASK;

    if (maxExclusive <= minInclusive)
        return (PHYSICAL_ADDRESS)0;

    uintptr_t flags = KeAcquireSpinlock(&gFrameStack.Lock);

    // Search high-to-low within the requested window.  For low-memory users
    // such as the AP trampoline this naturally prefers the traditional
    // 0x70000-0x9f000 area and avoids consuming lower real-mode space first.
    for (size_t i = gFrameStack.Top; i > 0; i--) {
        size_t idx = i - 1;
        PHYSICAL_ADDRESS frame = gFrameStack.Stack[idx];
        if (frame >= minInclusive && frame < maxExclusive &&
            !(frame & (PAGE_SIZE - 1))) {
            frame = KiPopFrameAtIndexLocked(idx);
            KeReleaseSpinlock(&gFrameStack.Lock, flags);
            return frame;
        }
    }

    KeReleaseSpinlock(&gFrameStack.Lock, flags);
    return (PHYSICAL_ADDRESS)0;
}

PHYSICAL_ADDRESS MmAllocateFrameBelow(PHYSICAL_ADDRESS maxExclusive) {
    return MmAllocateFrameInRange(PAGE_SIZE, maxExclusive);
}

__declspec(dllexport) PHYSICAL_ADDRESS MmAllocateFrame(void) {
    uintptr_t flags = KeAcquireSpinlock(&gFrameStack.Lock);
    if (gFrameStack.Top == 0) {
        KeReleaseSpinlock(&gFrameStack.Lock, flags);
        return (PHYSICAL_ADDRESS)0;
    }
    PHYSICAL_ADDRESS frame = gFrameStack.Stack[--gFrameStack.Top];
    KeReleaseSpinlock(&gFrameStack.Lock, flags);
    return frame;
}

void MmFreeFrame(PHYSICAL_ADDRESS frame) {
    if (!frame || (frame & (PAGE_SIZE - 1))) return;
    uintptr_t flags = KeAcquireSpinlock(&gFrameStack.Lock);
    if (gFrameStack.Top < gFrameStack.Capacity)
        gFrameStack.Stack[gFrameStack.Top++] = frame;
    // else: stack full - leak rather than corrupt
    KeReleaseSpinlock(&gFrameStack.Lock, flags);
}

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

    uintptr_t flags = KeAcquireSpinlock(&gFrameStack.Lock);

    if (gFrameStack.Top < count) {
        KeReleaseSpinlock(&gFrameStack.Lock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    // Frames are pushed in ascending address order by MmInitializePhysical,
    // so physically consecutive pages sit at consecutive stack indices.
    // Scan from bottom upward for a run of 'count' adjacent frames.
    size_t run = (size_t)-1;
    for (size_t i = 0; i + count <= gFrameStack.Top; i++) {
        int ok = 1;
        for (size_t j = 1; j < count; j++) {
            if (gFrameStack.Stack[i + j] != gFrameStack.Stack[i] + j * PAGE_SIZE) {
                ok = 0;
                // Skip past this broken run; the next candidate can't start
                // before the frame that broke the sequence.
                i += j - 1;  // outer loop adds 1 more
                break;
            }
        }
        if (ok) { run = i; break; }
    }

    if (run == (size_t)-1) {
        KeReleaseSpinlock(&gFrameStack.Lock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    PHYSICAL_ADDRESS base = gFrameStack.Stack[run];

    // Splice out the 'count' entries by sliding the tail down.
    size_t tail = gFrameStack.Top - run - count;
    for (size_t i = 0; i < tail; i++)
        gFrameStack.Stack[run + i] = gFrameStack.Stack[run + count + i];
    gFrameStack.Top -= count;

    KeReleaseSpinlock(&gFrameStack.Lock, flags);
    return base;
}

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

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