#include <VM/PMM.h>
#include <Lib/String.h>

extern char _kernelStart[];
extern char _kernelEnd[];

static inline Size BitmapGetByteIndex(Address address) {
    return (address / kVMPageSize) / kVMBlocksPerByte;
}

static inline UInt8 BitmapGetBitOffset(Address address) {
    return (UInt8)((address / kVMPageSize) % kVMBlocksPerByte);
}

static inline Boolean BitmapTest(const MemoryPointer bitmap, Address address) {
    return (bitmap[BitmapGetByteIndex(address)] & (1U << BitmapGetBitOffset(address))) != 0;
}

static inline void BitmapSet(MemoryPointer bitmap, Address address) {
    bitmap[BitmapGetByteIndex(address)] |= (1U << BitmapGetBitOffset(address));
}

static inline void BitmapUnset(MemoryPointer bitmap, Address address) {
    bitmap[BitmapGetByteIndex(address)] &= ~(1U << BitmapGetBitOffset(address));
}

static MemoryPointer sPMMBitmap;
static Size sPMMBitmapSize;
static Size sPMMTotalPages;

void PMMInitialize(VMBootMemoryMap* bootMap) {
    UInt32 vIndex = bootMap->reservedCount;
    bootMap->reserved[vIndex].base = 0x0;
    bootMap->reserved[vIndex].size = bootMap->totalRAM.base; 
    bootMap->reservedCount++;
    
    UInt32 kIndex = bootMap->reservedCount;
    bootMap->reserved[kIndex].base = (Address)_kernelStart;
    bootMap->reserved[kIndex].size = (Address)_kernelEnd - (Address)_kernelStart;
    bootMap->reservedCount++;

    sPMMTotalPages = bootMap->totalRAM.size / kVMPageSize;
    sPMMBitmapSize = sPMMTotalPages / kVMBlocksPerByte;
    sPMMBitmap = (MemoryPointer)_kernelEnd;

    StringSet(sPMMBitmap, 0, sPMMBitmapSize);

    UInt32 bIndex = bootMap->reservedCount;
    bootMap->reserved[bIndex].base = (Address)sPMMBitmap;
    bootMap->reserved[bIndex].size = sPMMBitmapSize;
    bootMap->reservedCount++;

    for (Size i = 0; i < bootMap->reservedCount; i++) {
        Address regionBase = bootMap->reserved[i].base;
        Size regionSize = bootMap->reserved[i].size;

        Size pagesToReserve = (regionSize + kVMPageSize - 1) / kVMPageSize;

        for (Size p = 0; p < pagesToReserve; p++) {
            Address pageAdress = regionBase + (p * kVMPageSize);
            BitmapSet(sPMMBitmap, pageAdress);
        }
    }
}

Pointer PMMAllocatePage() {
    for (Size i = 0; i < sPMMBitmapSize; i++) {
        if (sPMMBitmap[i] == 0xFF) continue;
        for (Size bit = 0; bit < kVMBlocksPerByte; bit++) {
            if ((sPMMBitmap[i] & (1 << bit)) == 0) {
                Address address = (i * kVMBlocksPerByte + bit) * kVMPageSize;
                BitmapSet(sPMMBitmap, address);
                return (Pointer)address;
            }
        }
    }

    return nullptr;
}

void PMMFreePage(Address address) {
    BitmapUnset(sPMMBitmap, address);
}