#include <string.h>
#include <stdio.h>

#include <mm/pmm.h>


#define MAX_PAGES (1024 * 1024) /* 4GB / 4KB */
static uint8_t bitmap[MAX_PAGES / 8] __attribute__((section(".pmm_bitmap")));
static size_t total_pages;

#define BITMAP_SET(i)   (bitmap[(i) / 8] |=  (1 << ((i) % 8)))
#define BITMAP_CLEAR(i) (bitmap[(i) / 8] &= ~(1 << ((i) % 8)))
#define BITMAP_TEST(i)  (bitmap[(i) / 8] &   (1 << ((i) % 8)))

void pmm_init(multiboot_info_t* mb)
{
#ifdef _DEBUG
  printf("[ PMM ] Initializing physical memory manager...\n");
#endif

  total_pages = MAX_PAGES;
  
  for (uint32_t i = 0; i < (total_pages / 8); i++)
  {
    bitmap[i] = 0xFF;
  }

  multiboot_memory_map_t* mmap = (void*)(uintptr_t)mb->mmap_addr;
  size_t entries = mb->mmap_length / sizeof(multiboot_memory_map_t);

  for (size_t i = 0; i < entries; i++)
  {
    if (mmap[i].type == 1) /* usable */
    {
      uint64_t start = mmap[i].addr;
      uint64_t end   = start + mmap[i].len;
      for (uint64_t addr = start; addr < end; addr += 0x1000)
      {
        if (addr >= 0x100000) /* skip below 1MB */
        {
          size_t idx = addr / 0x1000;
          BITMAP_CLEAR(idx);
        }
      }
    }
  }
  total_pages = MAX_PAGES;
  
#ifdef _DEBUG
  printf("[ PMM ] Physical memory manager initialized\n");
#endif
}

void* pmm_alloc_page(void)
{
  for (uint32_t i = 0; i < total_pages; ++i)
  {
    if (!BITMAP_TEST(i))
    {
      BITMAP_SET(i);
      return (void*)(i * 4096);
    }
  }
  
  printf("pmm_alloc_page(): No free page found!\n");
  
  return NULL;
}


void pmm_free_page(void* addr)
{
  size_t idx = (uintptr_t)addr / 0x1000;
  BITMAP_CLEAR(idx);
}