// mm_phys.c - physical memory manager
//
// Bitmap allocator over 4 KB physical frames.  One bitmap records which PFNs
// exist in usable RAM, and one bitmap records which usable PFNs are currently
// free.  This keeps single-frame allocation simple while making contiguous
// allocation and memory visualization deterministic.
#include <kernel/mm_phys.h>
#include <kernel/ke_spinlock.h>
#include <string.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 bitmap
// ---------------------------------------------------------------------------

typedef struct {
    uint8_t* Usable;       // bit = 1 when PFN belongs to usable RAM
    uint8_t* Free;         // bit = 1 when PFN is currently free
    size_t   BitmapBytes;
    size_t   MaxPfn;
    size_t   TotalFrames;  // usable frames available to the allocator
    size_t   FreeFrames;
    SPINLOCK Lock;
} FRAME_BITMAP;

static FRAME_BITMAP gFrameBitmap;

static inline PHYSICAL_ADDRESS KiAlignUp(PHYSICAL_ADDRESS value) {
    return (value + PAGE_SIZE - 1) & PAGE_MASK;
}

static inline PHYSICAL_ADDRESS KiAlignDown(PHYSICAL_ADDRESS value) {
    return value & PAGE_MASK;
}

static inline int KiBitTest(const uint8_t* bits, size_t index) {
    return (bits[index >> 3] & (uint8_t)(1u << (index & 7))) != 0;
}

static inline void KiBitSet(uint8_t* bits, size_t index) {
    bits[index >> 3] |= (uint8_t)(1u << (index & 7));
}

static inline void KiBitClear(uint8_t* bits, size_t index) {
    bits[index >> 3] &= (uint8_t)~(1u << (index & 7));
}

static inline int KiPfnUsableLocked(size_t pfn) {
    return pfn < gFrameBitmap.MaxPfn && KiBitTest(gFrameBitmap.Usable, pfn);
}

static inline int KiPfnFreeLocked(size_t pfn) {
    return KiPfnUsableLocked(pfn) && KiBitTest(gFrameBitmap.Free, pfn);
}

static void KiMarkUsableFreeRange(size_t firstPfn, size_t count) {
    for (size_t i = 0; i < count; i++) {
        size_t pfn = firstPfn + i;
        if (pfn >= gFrameBitmap.MaxPfn) break;
        if (!KiBitTest(gFrameBitmap.Usable, pfn)) {
            KiBitSet(gFrameBitmap.Usable, pfn);
            gFrameBitmap.TotalFrames++;
        }
        if (!KiBitTest(gFrameBitmap.Free, pfn)) {
            KiBitSet(gFrameBitmap.Free, pfn);
            gFrameBitmap.FreeFrames++;
        }
    }
}

static void KiRemoveUsablePfn(size_t pfn) {
    if (pfn >= gFrameBitmap.MaxPfn) return;
    if (KiBitTest(gFrameBitmap.Free, pfn)) {
        KiBitClear(gFrameBitmap.Free, pfn);
        if (gFrameBitmap.FreeFrames) gFrameBitmap.FreeFrames--;
    }
    if (KiBitTest(gFrameBitmap.Usable, pfn)) {
        KiBitClear(gFrameBitmap.Usable, pfn);
        if (gFrameBitmap.TotalFrames) gFrameBitmap.TotalFrames--;
    }
}

// ---------------------------------------------------------------------------
// Initialisation
// ---------------------------------------------------------------------------
bool MmInitializePhysical(const MEMORY_MAP* map, uint64_t physicalWindowBase) {
    if (!map) return false;
    gPhysBase = physicalWindowBase;

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

        PHYSICAL_ADDRESS base = KiAlignUp(r->Base);
        PHYSICAL_ADDRESS end = r->Base + r->Length;
        if (end < r->Base) end = (PHYSICAL_ADDRESS)-1;
        end = KiAlignDown(end);
        if (end <= base) continue;
        if (end > maxEnd) maxEnd = end;
    }

    size_t maxPfn = (size_t)(maxEnd / PAGE_SIZE);
    if (!maxPfn) return false;

    size_t bitmapBytes = (maxPfn + 7) / 8;
    size_t metadataBytes = bitmapBytes * 2;
    size_t metadataPages = (metadataBytes + PAGE_SIZE - 1) / PAGE_SIZE;
    PHYSICAL_ADDRESS metadataSize = (PHYSICAL_ADDRESS)metadataPages * PAGE_SIZE;

    // Keep allocator metadata out of low memory.  The bitmap is much smaller
    // than the old free-frame stack, so a first-fit search can otherwise choose
    // a conventional-memory hole below 1 MiB.  That range is still useful for
    // AP startup/legacy scratch, and stealing it makes early boot failures
    // unnecessarily silent because PMM initializes before normal debug output.
    const PHYSICAL_ADDRESS metadataMin = 0x100000ULL;
    PHYSICAL_ADDRESS metadataPhys = 0;
    int metadataFound = 0;
    for (uint32_t i = 0; i < map->RegionCount; i++) {
        const MEMORY_REGION* r = &map->Regions[i];
        if (r->Type != MemoryUsable) continue;

        PHYSICAL_ADDRESS base = KiAlignUp(r->Base);
        if (base < metadataMin) base = metadataMin;
        PHYSICAL_ADDRESS end = r->Base + r->Length;
        if (end < r->Base) end = (PHYSICAL_ADDRESS)-1;
        end = KiAlignDown(end);
        if (end <= base) continue;
        if (end - base >= metadataSize) {
            metadataPhys = base;
            metadataFound = 1;
            break;
        }
    }

    if (!metadataFound) return false;

    gFrameBitmap.BitmapBytes = bitmapBytes;
    gFrameBitmap.MaxPfn = maxPfn;
    gFrameBitmap.TotalFrames = 0;
    gFrameBitmap.FreeFrames = 0;
    gFrameBitmap.Lock = (SPINLOCK)SPINLOCK_INIT;
    gFrameBitmap.Usable = (uint8_t*)KiPhysToVirt(metadataPhys);
    gFrameBitmap.Free = gFrameBitmap.Usable + bitmapBytes;

    memset(gFrameBitmap.Usable, 0, metadataBytes);

    // Mark page-aligned usable RAM as free.  Unaligned tails are ignored rather
    // than handing malformed frame addresses to the rest of the kernel.
    for (uint32_t i = 0; i < map->RegionCount; i++) {
        const MEMORY_REGION* r = &map->Regions[i];
        if (r->Type != MemoryUsable) continue;

        PHYSICAL_ADDRESS base = KiAlignUp(r->Base);
        PHYSICAL_ADDRESS end = r->Base + r->Length;
        if (end < r->Base) end = (PHYSICAL_ADDRESS)-1;
        end = KiAlignDown(end);
        if (end <= base) continue;

        size_t firstPfn = (size_t)(base / PAGE_SIZE);
        size_t count = (size_t)((end - base) / PAGE_SIZE);
        KiMarkUsableFreeRange(firstPfn, count);
    }

    // PFN 0 is never a valid allocation result because 0 is the failure value.
    KiRemoveUsablePfn(0);

    // Hide the PMM metadata pages from the allocator entirely.  They are not
    // merely allocated; they are structural storage that must not become free if
    // some caller accidentally passes their PFNs to MmFreeFrame().
    size_t metadataFirstPfn = (size_t)(metadataPhys / PAGE_SIZE);
    for (size_t i = 0; i < metadataPages; i++)
        KiRemoveUsablePfn(metadataFirstPfn + i);

    return true;
}

// ---------------------------------------------------------------------------
// Public API
// ---------------------------------------------------------------------------
PHYSICAL_ADDRESS MmAllocateFrameInRange(PHYSICAL_ADDRESS minInclusive,
                                        PHYSICAL_ADDRESS maxExclusive) {
    minInclusive = KiAlignUp(minInclusive);
    maxExclusive = KiAlignDown(maxExclusive);
    if (maxExclusive <= minInclusive) return (PHYSICAL_ADDRESS)0;

    size_t firstPfn = (size_t)(minInclusive / PAGE_SIZE);
    size_t limitPfn = (size_t)(maxExclusive / PAGE_SIZE);
    if (limitPfn > gFrameBitmap.MaxPfn) limitPfn = gFrameBitmap.MaxPfn;
    if (firstPfn >= limitPfn) return (PHYSICAL_ADDRESS)0;

    uintptr_t flags = KeAcquireSpinlock(&gFrameBitmap.Lock);

    for (size_t pfn = limitPfn; pfn > firstPfn; pfn--) {
        size_t candidate = pfn - 1;
        if (KiPfnFreeLocked(candidate)) {
            KiBitClear(gFrameBitmap.Free, candidate);
            if (gFrameBitmap.FreeFrames) gFrameBitmap.FreeFrames--;
            KeReleaseSpinlock(&gFrameBitmap.Lock, flags);
            return (PHYSICAL_ADDRESS)candidate * PAGE_SIZE;
        }
    }

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

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

__declspec(dllexport) PHYSICAL_ADDRESS MmAllocateFrame(void) {
    return MmAllocateFrameInRange(PAGE_SIZE, (PHYSICAL_ADDRESS)gFrameBitmap.MaxPfn * PAGE_SIZE);
}

void MmFreeFrame(PHYSICAL_ADDRESS frame) {
    if (!frame || (frame & (PAGE_SIZE - 1))) return;
    size_t pfn = (size_t)(frame / PAGE_SIZE);

    uintptr_t flags = KeAcquireSpinlock(&gFrameBitmap.Lock);
    if (KiPfnUsableLocked(pfn) && !KiBitTest(gFrameBitmap.Free, pfn)) {
        KiBitSet(gFrameBitmap.Free, pfn);
        gFrameBitmap.FreeFrames++;
    }
    KeReleaseSpinlock(&gFrameBitmap.Lock, flags);
}

__declspec(dllexport) PHYSICAL_ADDRESS MmAllocateContiguous(size_t count) {
    if (!count || count > gFrameBitmap.MaxPfn) return (PHYSICAL_ADDRESS)0;

    uintptr_t flags = KeAcquireSpinlock(&gFrameBitmap.Lock);

    size_t runStart = 0;
    size_t runLength = 0;
    for (size_t pfn = 1; pfn < gFrameBitmap.MaxPfn; pfn++) {
        if (KiPfnFreeLocked(pfn)) {
            if (!runLength) runStart = pfn;
            runLength++;
            if (runLength == count) {
                for (size_t i = 0; i < count; i++)
                    KiBitClear(gFrameBitmap.Free, runStart + i);
                if (gFrameBitmap.FreeFrames >= count) gFrameBitmap.FreeFrames -= count;
                else gFrameBitmap.FreeFrames = 0;
                KeReleaseSpinlock(&gFrameBitmap.Lock, flags);
                return (PHYSICAL_ADDRESS)runStart * PAGE_SIZE;
            }
        } else {
            runLength = 0;
        }
    }

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

size_t MmFreeFrameCount(void) {
    uintptr_t flags = KeAcquireSpinlock(&gFrameBitmap.Lock);
    size_t count = gFrameBitmap.FreeFrames;
    KeReleaseSpinlock(&gFrameBitmap.Lock, flags);
    return count;
}

size_t MmTotalFrameCount(void) {
    return gFrameBitmap.TotalFrames;
}
void MmQueryPhysicalStats(MM_PHYSICAL_STATS* out) {
    if (!out) return;

    MM_PHYSICAL_STATS stats;
    memset(&stats, 0, sizeof(stats));

    uintptr_t flags = KeAcquireSpinlock(&gFrameBitmap.Lock);

    stats.TotalFrames = (uint64_t)gFrameBitmap.TotalFrames;
    stats.FreeFrames = (uint64_t)gFrameBitmap.FreeFrames;
    stats.UsedFrames = stats.TotalFrames - stats.FreeFrames;

    uint64_t run = 0;
    for (size_t pfn = 1; pfn < gFrameBitmap.MaxPfn; pfn++) {
        if (KiPfnFreeLocked(pfn)) {
            run++;
            continue;
        }

        if (run) {
            stats.FreeRunCount++;
            if (run > stats.LargestFreeRun)
                stats.LargestFreeRun = run;
            if (run == 1) stats.FreeRunHistogram[0]++;
            else if (run <= 3) stats.FreeRunHistogram[1]++;
            else if (run <= 7) stats.FreeRunHistogram[2]++;
            else if (run <= 15) stats.FreeRunHistogram[3]++;
            else if (run <= 63) stats.FreeRunHistogram[4]++;
            else if (run <= 255) stats.FreeRunHistogram[5]++;
            else if (run <= 1023) stats.FreeRunHistogram[6]++;
            else stats.FreeRunHistogram[7]++;
            run = 0;
        }
    }

    if (run) {
        stats.FreeRunCount++;
        if (run > stats.LargestFreeRun)
            stats.LargestFreeRun = run;
        if (run == 1) stats.FreeRunHistogram[0]++;
        else if (run <= 3) stats.FreeRunHistogram[1]++;
        else if (run <= 7) stats.FreeRunHistogram[2]++;
        else if (run <= 15) stats.FreeRunHistogram[3]++;
        else if (run <= 63) stats.FreeRunHistogram[4]++;
        else if (run <= 255) stats.FreeRunHistogram[5]++;
        else if (run <= 1023) stats.FreeRunHistogram[6]++;
        else stats.FreeRunHistogram[7]++;
    }

    KeReleaseSpinlock(&gFrameBitmap.Lock, flags);
    *out = stats;
}