#include <VM/PMM.h>
#include <Lib/String.h>
#include "../Common/bootinfo.h"

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

static Address sPMMRamBase = 0;
static MemoryPointer sPMMBitmap;
static Size sPMMBitmapSize;
static Size sPMMTotalPages;

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

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

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));
}

void PMMInitialize(VMBootMemoryMap* bootMap, Bootinfo* info) {
    sPMMRamBase = bootMap->totalRAM.base;
    sPMMTotalPages = bootMap->totalRAM.size / kVMPageSize;
    sPMMBitmapSize = sPMMTotalPages / kVMBlocksPerByte;
    sPMMBitmap = (MemoryPointer)_kernelEnd;
    MemorySet(sPMMBitmap, 0, sPMMBitmapSize);

    UInt32 kIndex = bootMap->reservedCount;
    bootMap->reserved[kIndex].base = (Address)_kernelStart;
    bootMap->reserved[kIndex].size = (Address)_kernelEnd - (Address)_kernelStart;
    bootMap->reservedCount++;

    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);
            if (pageAdress >= sPMMRamBase && pageAdress < (sPMMRamBase + bootMap->totalRAM.size)) {
                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 = sPMMRamBase + (i * kVMBlocksPerByte + bit) * kVMPageSize;
                BitmapSet(sPMMBitmap, address);
                return (Pointer)address;
            }
        }
    }

    return nullptr;
}

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