#include "pmm.h"
#include "../../kernel.h"
#include "../../mem.h"
#include <stdbool.h>
#include <stddef.h>

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

typedef struct
{
    uint64_t base_address;
    uint64_t length;
} memory_region_t;

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

void init_pmm(multiboot_info_t* mbd)
{
    memory_region_t usable_regions[USABLE_REGION_SIZE];
    size_t num_regions = 0;

    if (!(mbd->flags & (1 << 6)))
    {
        printf("Invalid memory map given by bootloader\n");
        panic();
    }

    uint64_t offset = 0;
    while (offset < mbd->mmap_length)
    {
        multiboot_memory_map_t* mmmt = (multiboot_memory_map_t*)(mbd->mmap_addr + offset);

        if (mmmt->type == MULTIBOOT_MEMORY_AVAILABLE && num_regions < USABLE_REGION_SIZE)
        {
            usable_regions[num_regions] = (memory_region_t){
                .base_address = mmmt->addr,
                .length = mmmt->len,
            };

            num_regions++;

            printf("Available memory: 0x%x - 0x%x (%u KB)\n", mmmt->addr, mmmt->addr + mmmt->len, mmmt->len / 1024);
        }

        offset += mmmt->size + sizeof(mmmt->size);
    }

    memset(page_bitmap, 0xFF, BITMAP_SIZE);

    for (size_t i = 0; i < num_regions; i++)
    {
        memory_region_t region = usable_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--;
            }
        }
    }

    printf("PMM initialized\n");
}

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

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