#include "alloc.h"
#include "common.h"
#include "print.h"
#include "../include/limine.h"


__attribute__((used, section(".limine_requests")))
static volatile struct limine_memmap_request mem_req = {
  .id = LIMINE_MEMMAP_REQUEST_ID,
  .revision = 0,
};

extern u32 fg, bg;

static void print_memmap_entry(struct limine_memmap_entry* entry) {
  u32 savefg = fg;
  u32 savebg = bg;
  static const char* type_strs[] = {
    "Usable",
    "Reserved",
    "ACPI Reclaimable",
    "ACPI NVS",
    "BAD!!!",
    "Bootloader Reclaimable",
    "Executable and Modules",
    "Framebuffer",
    "ACPI Tables",
  };

  static u32 type_colors[] = {
    0x00AF00,
    0xAF1F00,
    0x7f7f00,
    0x7f7f00,
    0xFF0000,
    0x007f7f,
    0x408000,
    0x0000FF,
    0x00FF00
  };

  bg = type_colors[entry->type];
  fg = ~bg;
  
  u64 addr = entry->base;
  u64 end = addr + entry->length;
  print("[");
  print(type_strs[entry->type]);
  print("]");
  print64(addr);
  print(" - ");
  print64(end);
  print("\n");
  fg = savefg;
  bg = savebg;
}



static size_t total_pages;
static u8* alloc_start;
static u8* alloc_end;
static u8* alloc_next;
static u64* alloc_bitmap;

void alloc_init() {
  print("Initializing Page Allocator\n");

  if(!mem_req.response) {
    print("Failed to get memory map :(\n)");
    hang();
  }
  
  u64 largest_base = 0;
  size_t largest_size = 0;

  for(size_t i = 0; i < mem_req.response->entry_count; i++) {
    struct limine_memmap_entry* entry = mem_req.response->entries[i];
    print_memmap_entry(entry);
    
    if(entry->type == LIMINE_MEMMAP_USABLE) {
      if(entry->length > largest_size) {
        largest_base = entry->base;
        largest_size = entry->length;
        print("\t\\LARGEST USABLE: ");
        print64(entry->length >> 20);
        print("MiB\n");
      }
    }
  }

  if(!largest_size) {
    print("Couldn't find a free page!!!\n");
    hang();
  }


  total_pages = largest_size / PAGE_SIZE;
  size_t usage_bitmap_size = (total_pages + 7) / 8;
  total_pages -= (usage_bitmap_size + PAGE_SIZE-1) / PAGE_SIZE;
  alloc_bitmap = (u64*)(phys_to_virt(largest_base));
  alloc_start = (u8*)(phys_to_virt(largest_base + usage_bitmap_size));
  alloc_next = alloc_start;
  alloc_end = (u8*)(phys_to_virt(largest_base + largest_size));
}

static size_t page(void* addr) {
  return ((u8*)addr - alloc_start) / PAGE_SIZE;
}

static void* addr(size_t page) {
  return alloc_start + page*PAGE_SIZE;
}

static void usepage(size_t page) {
   size_t idx = page / 64;
   size_t bit = page & 63;
   alloc_bitmap[idx] |= (1 << bit);
}

static void usepagerange(size_t page, size_t len) {
  for(size_t i = 0; i < len; len++) {
    usepage(i + page);
  }
}

static void unusepage(size_t page) {
   size_t idx = page / 64;
   size_t bit = page & 63;
   alloc_bitmap[idx] &= ~(1 << bit);
}

static void unusepagerange(size_t page, size_t len) {
  for(size_t i = 0; i < len; len++) {
    unusepage(i+page);
  }
}

static u8 getpageusage(size_t page) {
  return (alloc_bitmap[page / 8] >> (page & 7)) & 1;
}

void dealloc(void* addr, size_t size) {
  // adjust to page size
  size_t pages = (size+PAGE_SIZE-1) / PAGE_SIZE;
  unusepagerange(page(addr), pages);
}

void* alloc(size_t size) {
  size_t pages = (size + PAGE_SIZE - 1) / PAGE_SIZE;
  size_t page = 0;

  if(pages == 0 || pages >= total_pages) {
    return NULL;
  }

  size_t consecutive = 0;

  for(size_t i = 0; i < pages; i++) {
    u64 word = alloc_bitmap[i];
    if(word == 0) {
      if(!consecutive)
        page = i * 64;
      consecutive += 64;
      if(consecutive >= pages) {
        usepagerange(page, pages);
        return addr(page);
      }
      continue;
    }

    if(word == ~0ULL) {
      consecutive = 0;
      page = -1;
    }

    for(size_t bit = 0; bit < 64; bit++) {
      if((word & (1ULL << bit)) == 0) {
        if(!consecutive)
          page = i * 64 + bit;

        if(++consecutive >= pages) {
          usepagerange(page, pages);
          return addr(pages);
        }
      } else {
        consecutive = 0;
        page = -1;
      }
    }
  }


  return NULL;
}