#include "pmm.h"
#include "string.h"

static uint32_t* pmm_bitmap = NULL;
static uint32_t pmm_bitmap_size = 0;
static uint32_t total_pages = 0;

static pmm_stats_t global_stats;

static inline void bitmap_set(uint32_t page_index) {
    pmm_bitmap[page_index / 32] |= (1 << (page_index % 32));
}

static inline void bitmap_clear(uint32_t page_index) {
    pmm_bitmap[page_index / 32] &= ~(1 << (page_index % 32));
}

static inline int bitmap_test(uint32_t page_index) {
    return (pmm_bitmap[page_index / 32] & (1 << (page_index % 32))) != 0;
}

void pmm_init(struct multiboot_info* mbi, uint32_t kernel_end) {
    uint32_t mem_size = 128 * 1024 * 1024;
    total_pages = mem_size / PAGE_SIZE;

    pmm_bitmap = (uint32_t*)0x1000000;
    pmm_bitmap_size = total_pages / 8;

    uint32_t bitmap_end = kernel_end + pmm_bitmap_size;

    memset(pmm_bitmap, 0xFF, pmm_bitmap_size);

    struct multiboot_mmap_entry* mmap = (struct multiboot_mmap_entry*)mbi->mmap_addr;

    while((uint32_t)mmap < mbi->mmap_addr + mbi->mmap_length) {
        if (mmap->type == MULTIBOOT_MEMORY_AVAILABLE) {
            uint32_t start_addr = (uint32_t)mmap->addr;
            uint32_t end_addr = start_addr + (uint32_t)mmap->len;

            for (uint32_t addr = start_addr; addr < end_addr; addr += PAGE_SIZE) {
                uint32_t page_idx = addr / PAGE_SIZE;
                if (page_idx < total_pages) {
                    bitmap_clear(page_idx);
                }
            }
        }
        mmap = (struct multiboot_mmap_entry*)((uint32_t)mmap + mmap->size + sizeof(mmap->size));
    }

    for (uint32_t addr = 0; addr < bitmap_end; addr += PAGE_SIZE) {
        bitmap_set(addr / PAGE_SIZE);
    }
}

void* pmm_alloc_page(void) {
    for (uint32_t i = 0; i < total_pages; i++) {
        if (!bitmap_test(i)) {
            bitmap_set(i);
            return (void*)(i * PAGE_SIZE);
        }
    }
    return NULL;
}

void pmm_free_page(void* addr) {
    uint32_t page_idx = ((uint32_t)addr) / PAGE_SIZE;
    if (page_idx < total_pages) {
        bitmap_clear(page_idx);
    }
}

void pmm_get_stats(void) {
    uint32_t used = 0;
    uint32_t free = 0;

    for (uint32_t i = 0; i < total_pages; i++) {
        if (bitmap_test(i)) {
            used++;
        } else {
            free++;
        }
    }

    global_stats.total_pages = total_pages;
    global_stats.used_pages = used;
    global_stats.free_pages = free;
    global_stats.total_memory = total_pages << 12;
    global_stats.used_memory = used << 12;
    global_stats.free_memory = free << 12;
}

pmm_stats_t* pmm_get_stats_ptr(void) {
    return &global_stats;
}
