#include "pmm.h"
#include "../arch/arch.h"
#include <stddef.h>
#include <stdio.h>
#include <string.h>

#define BITMAP_SIZE 32768 // Supports up to 1GB of RAM
#define USABLE_REGION_SIZE 32

static uint8_t page_bitmap[BITMAP_SIZE];
static uint64_t total_pages = 0;
static uint64_t free_pages = 0;

void pmm_init()
{
    memset(page_bitmap, 0xFF, BITMAP_SIZE);

    for (size_t i = 0; i < memory_map.num_regions; i++)
    {
        memory_region_t region = memory_map.regions[i];

        uint64_t start_page = region.base_address / PAGE_SIZE;
        uint64_t num_pages = region.length / PAGE_SIZE;

        for (uint64_t page = start_page; page < start_page + num_pages; page++)
        {
            if (page < BITMAP_SIZE * 8)
            {
                page_bitmap[page / 8] &= ~(1 << (page % 8));
                free_pages++;
                total_pages++;
            }
            else
            {
                printf("System has more ram than the bitmap allows!\n");
                break;
            }
        }
    }

    // Reserve first 64MB which is reserved by boot code
    for (uint64_t page = 0; page < (64 * 1024 * 1024) / PAGE_SIZE; page++)
    {
        if (page < BITMAP_SIZE * 8)
        {
            if (!(page_bitmap[page / 8] & (1 << (page % 8))))
            {
                page_bitmap[page / 8] |= (1 << (page % 8));
                free_pages--;
            }
        }
    }
}

// Returns the address of first free physical page
uint64_t pmm_alloc_page()
{
    for (size_t i = 0; i < BITMAP_SIZE; i++)
    {
        if (page_bitmap[i] != 0xFF)
        {
            for (int bit = 0; bit < 8; bit++)
            {
                if (!(page_bitmap[i] & (1 << bit)))
                {
                    page_bitmap[i] |= (1 << bit);
                    free_pages--;
                    return ((i * 8 + bit) * PAGE_SIZE);
                }
            }
        }
    }

    return 0;
}

// Frees the physical page at the specified address
void pmm_free_page(uint64_t addr)
{
    uint64_t page = addr / PAGE_SIZE;
    if (page < BITMAP_SIZE * 8)
    {
        if (page_bitmap[page / 8] & (1 << (page % 8)))
        {
            page_bitmap[page / 8] &= ~(1 << (page % 8));
            free_pages++;
        }
    }
}