#include <stddef.h>

#include <mm/pmm.h>

#define PAGE_SIZE 4096
#define BITMAP_SIZE (1024 * 1024) /* Supports up to 4GB RAM (1 bit per page) */

static uint8_t bitmap[BITMAP_SIZE / 8];
static uint32_t total_pages;
static uint32_t used_pages = 0;

static inline void set_bit(uint32_t idx)
{
  bitmap[idx / 8] |= (1 << (idx % 8));
}

static inline void clear_bit(uint32_t idx)
{
  bitmap[idx / 8] &= ~(1 << (idx % 8));
}

static inline int test_bit(uint32_t idx)
{
  return (bitmap[idx / 8] >> (idx % 8)) & 1;
}

void pmm_init(multiboot_info_t* mb_info)
{
  total_pages = 0x100000; /* 4GB / 4KB = 1M pages */
  for (uint32_t i = 0; i < total_pages / 8; i++)
  {
      bitmap[i] = 0xFF; /* Mark all as used */
  }

  multiboot_memory_map_t* mmap = (multiboot_memory_map_t*) mb_info->mmap_addr;
  while ((uint32_t)mmap < mb_info->mmap_addr + mb_info->mmap_length)
  {
    if (mmap->type == 1) /* Usable */
    {
      uint64_t base = mmap->addr;
      uint64_t len = mmap->len;
      for (uint64_t addr = base; addr < base + len; addr += PAGE_SIZE)
      {
        if (addr >= 0x210000) /* Skip first 2.1MB, or ≈ 2.06MiB */
        {
          uint32_t idx = addr / PAGE_SIZE;
          
          if (idx >= total_pages)
          {
            continue; /* skip entries above 4GB */
          }
          
          clear_bit(idx);
          used_pages--;
        }
      }
    }
    mmap = (multiboot_memory_map_t*)((uint32_t)mmap + mmap->size + sizeof(mmap->size));
  }
}

void* alloc_page(void) {
  for (uint32_t i = 0; i < total_pages; i++)
  {
    if (!test_bit(i))
    {
      set_bit(i);
      used_pages++;
      return (void*)(i * PAGE_SIZE);
    }
  }
  return NULL; /* Out of memory */
}

void free_page(void* ptr)
{
  uint32_t idx = (uint32_t)ptr / PAGE_SIZE;
  clear_bit(idx);
}