#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include "../../include/limine.h"

static uint8_t* bitmap;
static uint64_t bitmap_pages;
static uint64_t total_pages;
static uint64_t last_index = 0;

// Helper to set/clear bits in our map
void bitmap_set(uint64_t page) {
    bitmap[page / 8] |= (1 << (page % 8));
}

void bitmap_clear(uint64_t page) {
    bitmap[page / 8] &= ~(1 << (page % 8));
}

bool bitmap_test(uint64_t page) {
    return bitmap[page / 8] & (1 << (page % 8));
}

/* * pmm_init:
 * Finds the biggest chunk of RAM and puts the bitmap there.
 * Then it marks all usable RAM as "free."
 */
void pmm_init(struct limine_memmap_response* memmap) {
    uint64_t max_addr = 0;
    uint64_t biggest_chunk_addr = 0;
    uint64_t biggest_chunk_len = 0;

    for (uint64_t i = 0; i < memmap->entry_count; i++) {
        struct limine_memmap_entry* entry = memmap->entries[i];
        if (entry->type == LIMINE_MEMMAP_USABLE) {
            if (entry->base + entry->length > max_addr)
                max_addr = entry->base + entry->length;
            
            if (entry->length > biggest_chunk_len) {
                biggest_chunk_len = entry->length;
                biggest_chunk_addr = entry->base;
            }
        }
    }

    total_pages = max_addr / 4096;
    bitmap_pages = (total_pages / 8) / 4096 + 1;
    
    // We store the bitmap in the biggest usable chunk of RAM
    bitmap = (uint8_t*)(biggest_chunk_addr + 0xffff800000000000); 
    
    // Mark EVERYTHING as used (locked) first
    for (uint64_t i = 0; i < total_pages / 8; i++) bitmap[i] = 0xFF;

    // Now, unmap only the truly usable sections
    for (uint64_t i = 0; i < memmap->entry_count; i++) {
        struct limine_memmap_entry* entry = memmap->entries[i];
        if (entry->type == LIMINE_MEMMAP_USABLE) {
            for (uint64_t j = 0; j < entry->length; j += 4096) {
                bitmap_clear((entry->base + j) / 4096);
            }
        }
    }

    // Lock the pages the bitmap itself is using
    for (uint64_t i = 0; i < bitmap_pages * 4096; i += 4096) {
        bitmap_set((biggest_chunk_addr + i) / 4096);
    }
}

/* pmm_alloc: Find a free page of RAM */
void* pmm_alloc() {
    for (uint64_t i = last_index; i < total_pages; i++) {
        if (!bitmap_test(i)) {
            bitmap_set(i);
            last_index = i;
            return (void*)(i * 4096);
        }
    }
    return NULL; // Out of memory!
}