// Kernel swap manager.
// The first backend is a fixed-size pagefile accessed through VFS offset I/O.
// Kernel memory and filesystem code remain non-pageable, so this layer is safe
// to use before the VMM has direct block-extents for pageout.

#include <stdint.h>
#include <stddef.h>
#include <string.h>
#include <kernel/mm_swap.h>
#include <kernel/mm_virt.h>
#include <kernel/mm_pfn.h>
#include <kernel/io_file.h>
#include <kernel/ke_spinlock.h>
#include <kernel/ke_debug.h>

typedef struct MM_SWAP_STATE {
    SPINLOCK Lock;
    FILE_OBJECT* File;
    uint8_t* Bitmap;
    uint32_t SlotCount;
    uint32_t UsedSlots;
    uint32_t HighWaterSlots;
    int Enabled;
    MM_SWAP_STATS Stats;
} MM_SWAP_STATE;

static MM_SWAP_STATE gSwap = { .Lock = SPINLOCK_INIT };

// Returns whether a swap slot is marked in use.
static int KiSwapTestSlot(uint32_t slot) {
    return (gSwap.Bitmap[slot >> 3] & (uint8_t)(1u << (slot & 7))) != 0;
}

// Marks a swap slot as allocated or free.
static void KiSwapSetSlot(uint32_t slot, int used) {
    uint8_t mask = (uint8_t)(1u << (slot & 7));
    if (used)
        gSwap.Bitmap[slot >> 3] |= mask;
    else
        gSwap.Bitmap[slot >> 3] &= (uint8_t)~mask;
}

// Allocates one pagefile slot.
static KSTATUS KiSwapAllocateSlot(uint32_t* slotOut) {
    if (!slotOut) return STATUS_INVALID_PARAMETER;
    *slotOut = 0;

    uintptr_t flags = KeAcquireSpinlock(&gSwap.Lock);
    if (!gSwap.Enabled || !gSwap.Bitmap || !gSwap.SlotCount) {
        KeReleaseSpinlock(&gSwap.Lock, flags);
        return STATUS_NOT_SUPPORTED;
    }

    for (uint32_t slot = 0; slot < gSwap.SlotCount; ++slot) {
        if (!KiSwapTestSlot(slot)) {
            KiSwapSetSlot(slot, 1);
            gSwap.UsedSlots++;
            if (gSwap.UsedSlots > gSwap.HighWaterSlots)
                gSwap.HighWaterSlots = gSwap.UsedSlots;
            gSwap.Stats.SlotAllocations++;
            *slotOut = slot;
            KeReleaseSpinlock(&gSwap.Lock, flags);
            return STATUS_SUCCESS;
        }
    }

    gSwap.Stats.SlotAllocationFailures++;
    KeReleaseSpinlock(&gSwap.Lock, flags);
    return STATUS_NO_MEMORY;
}

// Initializes the fixed pagefile backend.
KSTATUS MmSwapInitialize(SYS_STRING_VIEW path, uint64_t sizeBytes) {
    if (!path.Data || !path.Length || sizeBytes < PAGE_SIZE)
        return STATUS_INVALID_PARAMETER;
    if (sizeBytes & (PAGE_SIZE - 1))
        sizeBytes = (sizeBytes + PAGE_SIZE - 1) & ~(uint64_t)(PAGE_SIZE - 1);

    uintptr_t flags = KeAcquireSpinlock(&gSwap.Lock);
    if (gSwap.Enabled) {
        KeReleaseSpinlock(&gSwap.Lock, flags);
        return STATUS_SUCCESS;
    }
    KeReleaseSpinlock(&gSwap.Lock, flags);

    FILE_OBJECT* file = (FILE_OBJECT*)0;
    KSTATUS s = IoOpenFile(path, &file);
    if (!KSUCCESS(s)) {
        s = IoCreateFile(path, &file);
        if (!KSUCCESS(s))
            return s;
    }

    // Swap slots are never read before MmSwapPageOut overwrites the whole
    // page, so the filesystem may allocate backing clusters without zeroing
    // them. Normal files still get zero-on-extend semantics.
    file->Flags |= IO_FILE_FLAG_NO_CACHE | IO_FILE_FLAG_SKIP_ZERO_EXTEND;

    s = IoTruncateFile(file, sizeBytes);
    if (!KSUCCESS(s)) {
        IoCloseFile(file);
        return s;
    }

    uint64_t slots64 = sizeBytes / PAGE_SIZE;
    if (!slots64 || slots64 > 0xFFFFFFFFull) {
        IoCloseFile(file);
        return STATUS_INVALID_PARAMETER;
    }

    uint32_t slotCount = (uint32_t)slots64;
    size_t bitmapBytes = (slotCount + 7u) / 8u;
    uint8_t* bitmap = (uint8_t*)MmAllocatePoolTag(bitmapBytes, MM_POOL_TAG('s', 'w', 'a', 'p'));
    if (!bitmap) {
        IoCloseFile(file);
        return STATUS_NO_MEMORY;
    }
    memset(bitmap, 0, bitmapBytes);

    flags = KeAcquireSpinlock(&gSwap.Lock);
    if (gSwap.Enabled) {
        KeReleaseSpinlock(&gSwap.Lock, flags);
        MmFreePool(bitmap);
        IoCloseFile(file);
        return STATUS_SUCCESS;
    }

    gSwap.File = file;
    gSwap.Bitmap = bitmap;
    gSwap.SlotCount = slotCount;
    gSwap.UsedSlots = 0;
    gSwap.HighWaterSlots = 0;
    gSwap.Enabled = 1;
    memset(&gSwap.Stats, 0, sizeof(gSwap.Stats));
    gSwap.Stats.Enabled = 1;
    gSwap.Stats.TotalSlots = slotCount;
    gSwap.Stats.FreeSlots = slotCount;
    KeReleaseSpinlock(&gSwap.Lock, flags);

    DbgPrint("swap: initialized %u slots (%llu KiB) at %v\n",
             slotCount,
             (unsigned long long)(sizeBytes / 1024),
             path);
    return STATUS_SUCCESS;
}

// Returns whether swap is available for pageout.
int MmSwapIsEnabled(void) {
    return gSwap.Enabled;
}

// Writes one physical page to a newly allocated swap slot.
KSTATUS MmSwapPageOut(PHYSICAL_ADDRESS frame, uint32_t* slotOut) {
    if (!slotOut || !frame)
        return STATUS_INVALID_PARAMETER;
    *slotOut = 0;
    if (!gSwap.Enabled || !gSwap.File)
        return STATUS_NOT_SUPPORTED;

    uint32_t slot = 0;
    KSTATUS s = KiSwapAllocateSlot(&slot);
    if (!KSUCCESS(s)) {
        gSwap.Stats.PageOutFailures++;
        return s;
    }

    void* page = MmPfnToVirtual(frame);
    if (!page) {
        MmSwapFreeSlot(slot);
        gSwap.Stats.PageOutFailures++;
        return STATUS_FAILURE;
    }

    size_t written = 0;
    s = IoWriteFileAt(gSwap.File, (uint64_t)slot * PAGE_SIZE, page, PAGE_SIZE, &written);
    if (!KSUCCESS(s) || written != PAGE_SIZE) {
        MmSwapFreeSlot(slot);
        gSwap.Stats.PageOutFailures++;
        return KSUCCESS(s) ? STATUS_DEVICE_ERROR : s;
    }

    gSwap.Stats.PageOuts++;
    *slotOut = slot;
    return STATUS_SUCCESS;
}

// Reads one swap slot into an already allocated physical page.
KSTATUS MmSwapPageIn(uint32_t slot, PHYSICAL_ADDRESS frame) {
    if (!frame)
        return STATUS_INVALID_PARAMETER;
    if (!gSwap.Enabled || !gSwap.File)
        return STATUS_NOT_SUPPORTED;

    uintptr_t flags = KeAcquireSpinlock(&gSwap.Lock);
    int valid = slot < gSwap.SlotCount && KiSwapTestSlot(slot);
    KeReleaseSpinlock(&gSwap.Lock, flags);
    if (!valid) {
        gSwap.Stats.PageInFailures++;
        return STATUS_NOT_FOUND;
    }

    void* page = MmPfnToVirtual(frame);
    if (!page) {
        gSwap.Stats.PageInFailures++;
        return STATUS_FAILURE;
    }

    size_t read = 0;
    KSTATUS s = IoReadFileAt(gSwap.File, (uint64_t)slot * PAGE_SIZE, page, PAGE_SIZE, &read);
    if (!KSUCCESS(s) || read != PAGE_SIZE) {
        gSwap.Stats.PageInFailures++;
        return KSUCCESS(s) ? STATUS_DEVICE_ERROR : s;
    }

    gSwap.Stats.PageIns++;
    return STATUS_SUCCESS;
}

// Releases a swap slot after the owning VAD page no longer references it.
void MmSwapFreeSlot(uint32_t slot) {
    uintptr_t flags = KeAcquireSpinlock(&gSwap.Lock);
    if (gSwap.Enabled && slot < gSwap.SlotCount && KiSwapTestSlot(slot)) {
        KiSwapSetSlot(slot, 0);
        if (gSwap.UsedSlots)
            gSwap.UsedSlots--;
        gSwap.Stats.SlotFrees++;
    }
    KeReleaseSpinlock(&gSwap.Lock, flags);
}

// Copies swap telemetry for user diagnostics.
void MmSwapQueryStats(MM_SWAP_STATS* out) {
    if (!out) return;
    uintptr_t flags = KeAcquireSpinlock(&gSwap.Lock);
    MM_SWAP_STATS stats = gSwap.Stats;
    stats.Enabled = gSwap.Enabled ? 1 : 0;
    stats.TotalSlots = gSwap.SlotCount;
    stats.UsedSlots = gSwap.UsedSlots;
    stats.FreeSlots = gSwap.SlotCount >= gSwap.UsedSlots ? gSwap.SlotCount - gSwap.UsedSlots : 0;
    stats.HighWaterSlots = gSwap.HighWaterSlots;
    KeReleaseSpinlock(&gSwap.Lock, flags);
    *out = stats;
}