// Physical frame database.
// This file consumes the boot memory map and builds one PFN descriptor for each
// managed physical frame. The first implementation still allocates by scanning,
// but callers now get typed, reference-counted frames instead of anonymous
// bitmap bits.

#include <stdbool.h>
#include <kernel/mm_pfn.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_irql.h>
#include <string.h>

extern KSTATUS MmVadTrimOnePageForPressure(void);

static uint64_t gPhysBase = 0;
static MM_PFN* gPfnDatabase = (MM_PFN*)0;
static size_t gMaxPfn = 0;
static size_t gTotalFrames = 0;
static size_t gFreeFrames = 0;
static SPINLOCK gPfnLock = SPINLOCK_INIT;
static uint64_t gPfnAllocations = 0;
static uint64_t gPfnFrees = 0;
static uint64_t gPfnRetains = 0;
static uint64_t gPfnReleases = 0;
static uint64_t gPfnPins = 0;
static uint64_t gPfnUnpins = 0;
static uint64_t gPfnFailedAllocations = 0;
static uint64_t gPfnFailedContiguousAllocations = 0;
static uint64_t gPfnFailedInvalidRange = 0;
static uint64_t gPfnFailedEmptyRange = 0;
static uint64_t gPfnFailedReserveDenied = 0;
static uint64_t gPfnFailedNoFrame = 0;
static uint64_t gPfnFailedContiguousInvalidRequest = 0;
static uint64_t gPfnFailedContiguousReserveDenied = 0;
static uint64_t gPfnFailedContiguousNoRun = 0;
static uint64_t gPfnReclaimAttempts = 0;
static uint64_t gPfnReclaimPagesRequested = 0;
static uint64_t gPfnReclaimPagesSucceeded = 0;
static uint64_t gPfnReclaimPagesFailed = 0;
static uint64_t gPfnReclaimReentrantSkips = 0;
static volatile int gPfnReclaiming = 0;

static PHYSICAL_ADDRESS KiClaimPfnLocked(size_t pfn, MM_PFN_USAGE usage);

#define MM_PFN_KERNEL_RESERVE_FRAMES 128u
#define MM_PFN_RECLAIM_LOW_WATER     512u
#define MM_PFN_RECLAIM_HIGH_WATER    768u

// Converts a physical address into the direct-map alias used during early memory setup.
static inline void* KiPhysToVirt(PHYSICAL_ADDRESS phys) {
    return (void*)(phys + gPhysBase);
}

// Rounds a physical address up to the next whole frame boundary.
static inline PHYSICAL_ADDRESS KiAlignUp(PHYSICAL_ADDRESS value) {
    return (value + PAGE_SIZE - 1) & PAGE_MASK;
}

// Rounds a physical address down to the previous whole frame boundary.
static inline PHYSICAL_ADDRESS KiAlignDown(PHYSICAL_ADDRESS value) {
    return value & PAGE_MASK;
}

// Returns whether a boot memory type describes RAM that the PFN database should track.
static bool KiMemoryTypeManaged(MEMORY_TYPE type) {
    return type == MemoryUsable || type == MemoryBootloader || type == MemoryKernel;
}

// Returns whether an allocation may consume the emergency kernel reserve.
static bool KiPfnUsageMayUseReserve(MM_PFN_USAGE usage) {
    return usage == MmPfnUseKernel ||
           usage == MmPfnUsePageTable ||
           usage == MmPfnUseDma ||
           usage == MmPfnUseBootloader ||
           usage == MmPfnUsePfnDatabase ||
           usage == MmPfnUseMmio;
}

// Converts allocation purpose into the initial PFN state.
static MM_PFN_STATE KiAllocatedStateForUsage(MM_PFN_USAGE usage) {
    (void)usage;
    return MmPfnStateActive;
}

// Marks a PFN as used by a permanent structure that must never be returned to the free list.
static void KiReservePfnLocked(size_t pfn, MM_PFN_USAGE usage) {
    if (pfn >= gMaxPfn)
        return;

    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
        if (gFreeFrames)
            gFreeFrames--;
    }

    if (entry->State == MmPfnStateInvalid)
        gTotalFrames++;

    entry->State = MmPfnStateActive;
    entry->Usage = usage;
    entry->RefCount = 1;
    entry->PinCount = 1;
    entry->Flags = 0;
    entry->MappedCount = 0;
    entry->Owner = (void*)0;
    entry->BackingPage = 0;
}

// Marks a page-aligned RAM range as managed by the PFN database.
static void KiRegisterManagedRange(PHYSICAL_ADDRESS base,
                                   PHYSICAL_ADDRESS end,
                                   MM_PFN_STATE state,
                                   MM_PFN_USAGE usage) {
    size_t firstPfn = (size_t)(base / PAGE_SIZE);
    size_t count = (size_t)((end - base) / PAGE_SIZE);

    for (size_t i = 0; i < count; i++) {
        size_t pfn = firstPfn + i;
        if (pfn >= gMaxPfn)
            break;

        MM_PFN* entry = &gPfnDatabase[pfn];
        if (entry->State == MmPfnStateInvalid)
            gTotalFrames++;

        if (entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
            if (gFreeFrames)
                gFreeFrames--;
        }

        entry->State = state;
        entry->Usage = usage;
        entry->RefCount = (state == MmPfnStateFree || state == MmPfnStateZeroed) ? 0 : 1;
        entry->PinCount = 0;
        entry->Flags = 0;
        entry->MappedCount = 0;
        entry->Owner = (void*)0;
        entry->BackingPage = 0;

        if (state == MmPfnStateFree || state == MmPfnStateZeroed)
            gFreeFrames++;
    }
}

// Returns whether the future page-replacement code may consider this frame a candidate.
static bool KiPfnEvictableLocked(const MM_PFN* entry) {
    if (!entry || entry->PinCount)
        return false;
    if (entry->State != MmPfnStateActive && entry->State != MmPfnStateStandby && entry->State != MmPfnStateModified)
        return false;
    return entry->Usage == MmPfnUseUser ||
           entry->Usage == MmPfnUseCleanFileCache ||
           entry->Usage == MmPfnUseImage ||
           entry->Usage == MmPfnUseMappedFile ||
           entry->Usage == MmPfnUseAnonymous ||
           entry->Usage == MmPfnUseCow;
}

// Returns whether a standby frame can be counted as immediately reclaimable.
static bool KiPfnReclaimableLocked(const MM_PFN* entry) {
    if (!entry || entry->PinCount || entry->MappedCount)
        return false;
    if (entry->State != MmPfnStateStandby || entry->RefCount != 1)
        return false;
    return entry->Usage == MmPfnUseCleanFileCache || entry->Usage == MmPfnUseImage;
}

// Returns true when a PFN can satisfy a new allocation.
static bool KiPfnAvailableLocked(size_t pfn) {
    if (pfn >= gMaxPfn)
        return false;
    MM_PFN* entry = &gPfnDatabase[pfn];
    return entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed;
}

// Tries to create breathing room by swapping private user pages before an allocation fails.
static void KiPfnReclaimForPressure(void) {
    if (!KeCanWait())
        return;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    if (gPfnReclaiming) {
        gPfnReclaimReentrantSkips++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return;
    }
    if (gFreeFrames >= MM_PFN_RECLAIM_LOW_WATER) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return;
    }

    gPfnReclaiming = 1;
    gPfnReclaimAttempts++;
    KeReleaseSpinlock(&gPfnLock, flags);

    for (uint32_t i = 0; i < 64; ++i) {
        flags = KeAcquireSpinlock(&gPfnLock);
        if (gFreeFrames >= MM_PFN_RECLAIM_HIGH_WATER) {
            KeReleaseSpinlock(&gPfnLock, flags);
            break;
        }
        gPfnReclaimPagesRequested++;
        KeReleaseSpinlock(&gPfnLock, flags);

        KSTATUS s = MmVadTrimOnePageForPressure();

        flags = KeAcquireSpinlock(&gPfnLock);
        if (KSUCCESS(s))
            gPfnReclaimPagesSucceeded++;
        else
            gPfnReclaimPagesFailed++;
        KeReleaseSpinlock(&gPfnLock, flags);

        if (!KSUCCESS(s))
            break;
    }

    flags = KeAcquireSpinlock(&gPfnLock);
    gPfnReclaiming = 0;
    KeReleaseSpinlock(&gPfnLock, flags);
}

// Claims a free PFN in the requested range while the PFN lock is held.
static PHYSICAL_ADDRESS KiClaimAvailableInRangeLocked(size_t firstPfn, size_t limitPfn, MM_PFN_USAGE usage) {
    for (size_t pfn = limitPfn; pfn > firstPfn; pfn--) {
        size_t candidate = pfn - 1;
        if (KiPfnAvailableLocked(candidate))
            return KiClaimPfnLocked(candidate, usage);
    }
    return (PHYSICAL_ADDRESS)0;
}

// Claims one free PFN for the requested usage.
static PHYSICAL_ADDRESS KiClaimPfnLocked(size_t pfn, MM_PFN_USAGE usage) {
    MM_PFN* entry = &gPfnDatabase[pfn];
    entry->State = KiAllocatedStateForUsage(usage);
    entry->Usage = usage;
    entry->RefCount = 1;
    entry->PinCount = (usage == MmPfnUsePageTable) ? 1 : 0;
    entry->Flags = 0;
    entry->MappedCount = 0;
    entry->Owner = (void*)0;
    entry->BackingPage = 0;
    if (gFreeFrames)
        gFreeFrames--;
    gPfnAllocations++;
    return MmPfnAddress(pfn);
}

// Initializes the PFN database from the bootloader memory map.
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 (!KiMemoryTypeManaged(r->Type))
            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;
    }

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

    size_t pfnBytes = gMaxPfn * sizeof(MM_PFN);
    size_t pfnPages = (pfnBytes + PAGE_SIZE - 1) / PAGE_SIZE;
    PHYSICAL_ADDRESS metadataSize = (PHYSICAL_ADDRESS)pfnPages * PAGE_SIZE;

    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;

    gPfnDatabase = (MM_PFN*)KiPhysToVirt(metadataPhys);
    memset(gPfnDatabase, 0, metadataSize);
    gTotalFrames = 0;
    gFreeFrames = 0;
    gPfnLock = (SPINLOCK)SPINLOCK_INIT;
    gPfnAllocations = 0;
    gPfnFrees = 0;
    gPfnRetains = 0;
    gPfnReleases = 0;
    gPfnPins = 0;
    gPfnUnpins = 0;
    gPfnFailedAllocations = 0;
    gPfnFailedContiguousAllocations = 0;
    gPfnFailedInvalidRange = 0;
    gPfnFailedEmptyRange = 0;
    gPfnFailedReserveDenied = 0;
    gPfnFailedNoFrame = 0;
    gPfnFailedContiguousInvalidRequest = 0;
    gPfnFailedContiguousReserveDenied = 0;
    gPfnFailedContiguousNoRun = 0;
    gPfnReclaimAttempts = 0;
    gPfnReclaimPagesRequested = 0;
    gPfnReclaimPagesSucceeded = 0;
    gPfnReclaimPagesFailed = 0;
    gPfnReclaimReentrantSkips = 0;

    for (uint32_t i = 0; i < map->RegionCount; i++) {
        const MEMORY_REGION* r = &map->Regions[i];
        if (!KiMemoryTypeManaged(r->Type))
            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 (r->Type == MemoryUsable) {
            KiRegisterManagedRange(base, end, MmPfnStateFree, MmPfnUseNone);
        } else if (r->Type == MemoryBootloader) {
            KiRegisterManagedRange(base, end, MmPfnStateActive, MmPfnUseBootloader);
        } else {
            KiRegisterManagedRange(base, end, MmPfnStateActive, MmPfnUseKernel);
            for (size_t pfn = (size_t)(base / PAGE_SIZE); pfn < (size_t)(end / PAGE_SIZE) && pfn < gMaxPfn; pfn++)
                gPfnDatabase[pfn].PinCount = 1;
        }
    }

    // PFN 0 remains reserved because a zero physical address is the allocation failure value.
    KiReservePfnLocked(0, MmPfnUseKernel);

    // The PFN database owns its storage permanently. It is structural metadata,
    // not an allocation that can later be freed through a stale physical address.
    size_t metadataFirstPfn = (size_t)(metadataPhys / PAGE_SIZE);
    for (size_t i = 0; i < pfnPages; i++)
        KiReservePfnLocked(metadataFirstPfn + i, MmPfnUsePfnDatabase);

    return true;
}

// Converts a physical frame address to a PFN index.
size_t MmPfnIndex(PHYSICAL_ADDRESS frame) {
    if (frame & (PAGE_SIZE - 1))
        return MM_PFN_INVALID_INDEX;
    size_t pfn = (size_t)(frame / PAGE_SIZE);
    if (pfn >= gMaxPfn)
        return MM_PFN_INVALID_INDEX;
    return pfn;
}

// Converts a PFN index to a physical frame address.
PHYSICAL_ADDRESS MmPfnAddress(size_t pfn) {
    return (PHYSICAL_ADDRESS)pfn * PAGE_SIZE;
}

// Returns the PFN descriptor for a page-aligned physical frame.
__declspec(dllexport) MM_PFN* MmPfnGet(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return (MM_PFN*)0;
    return &gPfnDatabase[pfn];
}

// Returns the direct-map virtual address for a managed physical frame.
__declspec(dllexport) void* MmPfnToVirtual(PHYSICAL_ADDRESS frame) {
    if (MmPfnIndex(frame) == MM_PFN_INVALID_INDEX)
        return (void*)0;
    return KiPhysToVirt(frame);
}

// Allocates one physical frame for the requested kernel subsystem.
__declspec(dllexport) PHYSICAL_ADDRESS MmPfnAllocate(MM_PFN_USAGE usage) {
    return MmPfnAllocateInRange(PAGE_SIZE, (PHYSICAL_ADDRESS)gMaxPfn * PAGE_SIZE, usage);
}

// Allocates one physical frame below a physical address limit.
__declspec(dllexport) PHYSICAL_ADDRESS MmPfnAllocateBelow(PHYSICAL_ADDRESS maxExclusive, MM_PFN_USAGE usage) {
    return MmPfnAllocateInRange(PAGE_SIZE, maxExclusive, usage);
}

// Allocates one physical frame inside a physical address range.
__declspec(dllexport) PHYSICAL_ADDRESS MmPfnAllocateInRange(PHYSICAL_ADDRESS minInclusive,
                                      PHYSICAL_ADDRESS maxExclusive,
                                      MM_PFN_USAGE usage) {
    minInclusive = KiAlignUp(minInclusive);
    maxExclusive = KiAlignDown(maxExclusive);
    if (maxExclusive <= minInclusive) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedAllocations++;
        gPfnFailedInvalidRange++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    size_t firstPfn = (size_t)(minInclusive / PAGE_SIZE);
    size_t limitPfn = (size_t)(maxExclusive / PAGE_SIZE);
    if (limitPfn > gMaxPfn)
        limitPfn = gMaxPfn;
    if (firstPfn >= limitPfn) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedAllocations++;
        gPfnFailedEmptyRange++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    KiPfnReclaimForPressure();

    int mayUseReserve = KiPfnUsageMayUseReserve(usage);

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    if (!mayUseReserve && gFreeFrames <= MM_PFN_KERNEL_RESERVE_FRAMES) {
        KeReleaseSpinlock(&gPfnLock, flags);
    } else {
        PHYSICAL_ADDRESS phys = KiClaimAvailableInRangeLocked(firstPfn, limitPfn, usage);
        if (phys) {
            KeReleaseSpinlock(&gPfnLock, flags);
            return phys;
        }
        KeReleaseSpinlock(&gPfnLock, flags);
    }

    KiPfnReclaimForPressure();

    flags = KeAcquireSpinlock(&gPfnLock);
    if (!mayUseReserve && gFreeFrames <= MM_PFN_KERNEL_RESERVE_FRAMES) {
        gPfnFailedAllocations++;
        gPfnFailedReserveDenied++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    PHYSICAL_ADDRESS phys = KiClaimAvailableInRangeLocked(firstPfn, limitPfn, usage);
    if (phys) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return phys;
    }
    gPfnFailedAllocations++;
    gPfnFailedNoFrame++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return (PHYSICAL_ADDRESS)0;
}

// Allocates a physically contiguous run inside an address and alignment window.
PHYSICAL_ADDRESS MmPfnAllocateContiguousInRange(size_t count,
                                                 PHYSICAL_ADDRESS minInclusive,
                                                 PHYSICAL_ADDRESS maxExclusive,
                                                 size_t alignment,
                                                 size_t boundary,
                                                 MM_PFN_USAGE usage) {
    if (!count || count > gMaxPfn) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousInvalidRequest++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    if (!alignment)
        alignment = PAGE_SIZE;
    if (alignment < PAGE_SIZE)
        alignment = PAGE_SIZE;
    if ((alignment & (alignment - 1)) != 0 ||
        (boundary && (boundary & (boundary - 1)) != 0)) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousInvalidRequest++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    if (count > ((size_t)-1) / PAGE_SIZE) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousInvalidRequest++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    size_t byteCount = count * PAGE_SIZE;
    if (boundary && byteCount > boundary) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousInvalidRequest++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    minInclusive = KiAlignUp(minInclusive);
    maxExclusive = KiAlignDown(maxExclusive);
    PHYSICAL_ADDRESS managedLimit = (PHYSICAL_ADDRESS)gMaxPfn * PAGE_SIZE;
    if (maxExclusive > managedLimit)
        maxExclusive = managedLimit;
    if (maxExclusive <= minInclusive ||
        (PHYSICAL_ADDRESS)byteCount > maxExclusive - minInclusive) {
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousInvalidRequest++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    if (!KiPfnUsageMayUseReserve(usage))
        KiPfnReclaimForPressure();

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    if (!KiPfnUsageMayUseReserve(usage) && gFreeFrames < MM_PFN_KERNEL_RESERVE_FRAMES + count) {
        gPfnFailedContiguousAllocations++;
        gPfnFailedContiguousReserveDenied++;
        KeReleaseSpinlock(&gPfnLock, flags);
        return (PHYSICAL_ADDRESS)0;
    }

    size_t firstPfn = (size_t)(minInclusive / PAGE_SIZE);
    size_t limitPfn = (size_t)(maxExclusive / PAGE_SIZE);
    for (size_t runStart = firstPfn; runStart + count <= limitPfn; ++runStart) {
        PHYSICAL_ADDRESS runAddress = MmPfnAddress(runStart);
        if ((runAddress & (alignment - 1)) != 0)
            continue;
        if (boundary) {
            PHYSICAL_ADDRESS runEnd = runAddress + (PHYSICAL_ADDRESS)byteCount - 1;
            if ((runAddress & ~((PHYSICAL_ADDRESS)boundary - 1)) !=
                (runEnd & ~((PHYSICAL_ADDRESS)boundary - 1)))
                continue;
        }

        size_t available = 0;
        while (available < count && KiPfnAvailableLocked(runStart + available))
            available++;
        if (available != count) {
            runStart += available;
            continue;
        }

        for (size_t i = 0; i < count; ++i)
            KiClaimPfnLocked(runStart + i, usage);
        KeReleaseSpinlock(&gPfnLock, flags);
        return runAddress;
    }

    gPfnFailedContiguousAllocations++;
    gPfnFailedContiguousNoRun++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return (PHYSICAL_ADDRESS)0;
}

// Allocates a physically contiguous run for DMA or hardware structures.
__declspec(dllexport) PHYSICAL_ADDRESS MmPfnAllocateContiguous(size_t count, MM_PFN_USAGE usage) {
    return MmPfnAllocateContiguousInRange(count,
                                           PAGE_SIZE,
                                           (PHYSICAL_ADDRESS)gMaxPfn * PAGE_SIZE,
                                           PAGE_SIZE,
                                           0,
                                           usage);
}

// Increments the reference count on an allocated frame.
__declspec(dllexport) KSTATUS MmPfnRetain(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateInvalid || entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (entry->RefCount == UINT16_MAX) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_FAILURE;
    }
    entry->RefCount++;
    gPfnRetains++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Releases a reference and returns the frame to the free list when the last owner leaves.
__declspec(dllexport) KSTATUS MmPfnRelease(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateInvalid || entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (entry->PinCount) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_ACCESS_DENIED;
    }
    gPfnReleases++;
    if (entry->RefCount > 1) {
        entry->RefCount--;
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_SUCCESS;
    }
    if (entry->MappedCount) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_ACCESS_DENIED;
    }

    entry->State = MmPfnStateFree;
    entry->Usage = MmPfnUseNone;
    entry->RefCount = 0;
    entry->PinCount = 0;
    entry->Flags = 0;
    entry->MappedCount = 0;
    entry->Owner = (void*)0;
    entry->BackingPage = 0;
    gFreeFrames++;
    gPfnFrees++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Frees a frame that is expected to have a single owner.
__declspec(dllexport) KSTATUS MmPfnFree(PHYSICAL_ADDRESS frame) {
    return MmPfnRelease(frame);
}

// Pins a frame so later page replacement cannot evict it.
__declspec(dllexport) KSTATUS MmPfnPin(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateInvalid || entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (entry->PinCount == UINT16_MAX) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_FAILURE;
    }
    entry->PinCount++;
    gPfnPins++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Removes one pin from a frame.
__declspec(dllexport) KSTATUS MmPfnUnpin(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (!entry->PinCount) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    entry->PinCount--;
    gPfnUnpins++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Records that a VAD now maps this frame into a user address space.
__declspec(dllexport) KSTATUS MmPfnAddMapping(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateInvalid || entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    if (entry->MappedCount == UINT32_MAX) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_FAILURE;
    }
    entry->MappedCount++;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Removes one user VAD mapping reference from a frame.
__declspec(dllexport) KSTATUS MmPfnRemoveMapping(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State == MmPfnStateInvalid || entry->State == MmPfnStateFree || entry->State == MmPfnStateZeroed ||
        entry->MappedCount == 0) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }
    entry->MappedCount--;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Moves a clean, unshared frame to standby so pressure can reclaim it later.
__declspec(dllexport) KSTATUS MmPfnMoveToStandby(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State != MmPfnStateActive || entry->PinCount || entry->MappedCount || entry->RefCount != 1) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_ACCESS_DENIED;
    }
    if (entry->Usage != MmPfnUseCleanFileCache && entry->Usage != MmPfnUseImage) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    entry->State = MmPfnStateStandby;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Reactivates a standby frame after a cache or image lookup found it still useful.
__declspec(dllexport) KSTATUS MmPfnActivate(PHYSICAL_ADDRESS frame) {
    size_t pfn = MmPfnIndex(frame);
    if (pfn == MM_PFN_INVALID_INDEX)
        return NTLSTATUS_INVALID_PARAMETER;

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    MM_PFN* entry = &gPfnDatabase[pfn];
    if (entry->State != MmPfnStateStandby || !entry->RefCount) {
        KeReleaseSpinlock(&gPfnLock, flags);
        return NTLSTATUS_INVALID_PARAMETER;
    }

    entry->State = MmPfnStateActive;
    KeReleaseSpinlock(&gPfnLock, flags);
    return NTLSTATUS_SUCCESS;
}

// Returns the number of frames that are immediately available for allocation.
size_t MmFreeFrameCount(void) {
    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    size_t count = gFreeFrames;
    KeReleaseSpinlock(&gPfnLock, flags);
    return count;
}

// Returns the number of RAM frames managed by the PFN database.
size_t MmTotalFrameCount(void) {
    return gTotalFrames;
}

// Returns detailed PFN state counters for VMM self-tests and diagnostics.
void MmPfnQueryStats(MM_PFN_STATS* out) {
    if (!out)
        return;

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

    // Detailed statistics are diagnostic snapshots. Scan bounded PFN batches
    // so an ordinary query cannot monopolize the global allocator lock on
    // machines with millions of physical frames. Counters are sampled at end.
    const size_t pfnCount = gMaxPfn;
    for (size_t base = 0; base < pfnCount;) {
        size_t end = pfnCount - base < 1024u ? pfnCount : base + 1024u;
        uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
        for (size_t pfn = base; pfn < end; ++pfn) {
        const MM_PFN* entry = &gPfnDatabase[pfn];
        if (entry->State == MmPfnStateInvalid)
            continue;
        switch ((MM_PFN_STATE)entry->State) {
        case MmPfnStateReserved: stats.ReservedFrames++; break;
        case MmPfnStateZeroed:   stats.ZeroedFrames++; break;
        case MmPfnStateActive:   stats.ActiveFrames++; break;
        case MmPfnStateStandby:  stats.StandbyFrames++; break;
        case MmPfnStateModified: stats.ModifiedFrames++; break;
        default: break;
        }
        if (entry->PinCount)
            stats.PinnedFrames++;
        if (entry->MappedCount)
            stats.MappedFrames++;
        if (KiPfnEvictableLocked(entry))
            stats.EvictableFrames++;
        else if (entry->State != MmPfnStateFree && entry->State != MmPfnStateZeroed)
            stats.NonEvictableFrames++;
        if (KiPfnReclaimableLocked(entry))
            stats.ReclaimableFrames++;

        switch ((MM_PFN_USAGE)entry->Usage) {
        case MmPfnUseKernel:         stats.KernelFrames++; break;
        case MmPfnUsePageTable:      stats.PageTableFrames++; break;
        case MmPfnUseUser:           stats.UserFrames++; break;
        case MmPfnUseCleanFileCache:
            stats.CleanFileCacheFrames++;
            if (entry->State == MmPfnStateStandby) stats.StandbyCleanFileCacheFrames++;
            else if (entry->State == MmPfnStateActive) stats.ActiveCleanFileCacheFrames++;
            break;
        case MmPfnUseImage:
            stats.ImageFrames++;
            if (entry->State == MmPfnStateStandby) stats.StandbyImageFrames++;
            else if (entry->State == MmPfnStateActive) stats.ActiveImageFrames++;
            break;
        case MmPfnUseMappedFile:     stats.MappedFileFrames++; break;
        case MmPfnUseAnonymous:      stats.AnonymousFrames++; break;
        case MmPfnUseCow:            stats.CowFrames++; break;
        case MmPfnUseSurface:        stats.SurfaceFrames++; break;
        case MmPfnUseDriver:         stats.DriverFrames++; break;
        case MmPfnUseDma:            stats.DmaFrames++; break;
        case MmPfnUseBootloader:     stats.BootloaderFrames++; break;
        case MmPfnUsePfnDatabase:    stats.PfnDatabaseFrames++; break;
        case MmPfnUseMmio:           stats.MmioFrames++; break;
        default: break;
        }
        }
        KeReleaseSpinlock(&gPfnLock, flags);
        base = end;
    }

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    stats.TotalFrames = (uint64_t)gTotalFrames;
    stats.FreeFrames = (uint64_t)gFreeFrames;
    stats.UsedFrames = stats.TotalFrames - stats.FreeFrames;
    stats.AvailableFrames = stats.FreeFrames + stats.ReclaimableFrames;

    stats.AllocationCount = gPfnAllocations;
    stats.FreeCount = gPfnFrees;
    stats.RetainCount = gPfnRetains;
    stats.ReleaseCount = gPfnReleases;
    stats.PinCount = gPfnPins;
    stats.UnpinCount = gPfnUnpins;
    stats.FailedAllocations = gPfnFailedAllocations;
    stats.FailedContiguousAllocations = gPfnFailedContiguousAllocations;
    stats.FailedInvalidRange = gPfnFailedInvalidRange;
    stats.FailedEmptyRange = gPfnFailedEmptyRange;
    stats.FailedReserveDenied = gPfnFailedReserveDenied;
    stats.FailedNoFrame = gPfnFailedNoFrame;
    stats.FailedContiguousInvalidRequest = gPfnFailedContiguousInvalidRequest;
    stats.FailedContiguousReserveDenied = gPfnFailedContiguousReserveDenied;
    stats.FailedContiguousNoRun = gPfnFailedContiguousNoRun;
    stats.ReclaimAttempts = gPfnReclaimAttempts;
    stats.ReclaimPagesRequested = gPfnReclaimPagesRequested;
    stats.ReclaimPagesSucceeded = gPfnReclaimPagesSucceeded;
    stats.ReclaimPagesFailed = gPfnReclaimPagesFailed;
    stats.ReclaimReentrantSkips = gPfnReclaimReentrantSkips;

    KeReleaseSpinlock(&gPfnLock, flags);
    *out = stats;
}

// Returns the PMM counters expected by the existing stats syscall.
void MmQueryPhysicalStats(MM_PHYSICAL_STATS* out) {
    if (!out)
        return;

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

    uintptr_t flags = KeAcquireSpinlock(&gPfnLock);
    stats.TotalFrames = (uint64_t)gTotalFrames;
    stats.FreeFrames = (uint64_t)gFreeFrames;
    stats.UsedFrames = stats.TotalFrames - stats.FreeFrames;

    uint64_t run = 0;
    for (size_t pfn = 1; pfn < gMaxPfn; pfn++) {
        if (KiPfnAvailableLocked(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(&gPfnLock, flags);
    *out = stats;
}