#include <mm/pmm.h>
#include <mm/paging.h>
#include <string.h>
#include <stdio.h>
#include <mm_macros.h>

#include <mm/heap.h>

#define ALIGN4(x) (((x) + 3) & ~3)
#define MIN_BLOCK_SIZE 16

typedef struct block {
  size_t size;
  struct block* next;
  int free;
} block_t;

#define BLOCK_SIZE sizeof(block_t)

static uint8_t* heap_base;
static uint8_t* heap_end;
static size_t   heap_size;
static block_t* free_list;

void heap_init(uint32_t start, uint32_t size)
{
#ifdef _DEBUG
  printf("[ HEAP ] Initializing heap allocator...\n");
#endif

  heap_base = (uint8_t*) ALIGN4((uintptr_t) start);
  heap_end  = heap_base;
  heap_size = size;
  free_list = NULL;
  
  for (uint32_t i = 0; i < size; i += 0x1000)
  {
    map_page(pmm_alloc_page(), (void*)(start + i));
  }

  /* Set up initial free block */
  free_list = (block_t*)heap_base;
  free_list->size = heap_size - BLOCK_SIZE;
  free_list->free = 1;
  free_list->next = NULL;

  heap_end = heap_base + size;
  
#ifdef _DEBUG
  printf("[ HEAP ] Heap allocator initialized\n");
#endif
}

static block_t* find_free_block(size_t size)
{
  block_t* curr = free_list;
  while (curr)
  {
    if (curr->free && curr->size >= size)
    {
      return curr;
    }
    
    curr = curr->next;
  }
  
  printf("find_free_block(): No free block found!\n");
  return NULL;
}

static void split_block(block_t* blk, size_t size)
{
  if (blk->size >= size + BLOCK_SIZE + MIN_BLOCK_SIZE)
  {
    block_t* new_blk = (block_t*)((uint8_t*)blk + BLOCK_SIZE + size);
    new_blk->size = blk->size - size - BLOCK_SIZE;
    new_blk->free = 1;
    new_blk->next = blk->next;
    blk->next = new_blk;
    blk->size = size;
  }
}

void* malloc(size_t size)
{
  size = ALIGN4(size);
  block_t* blk = find_free_block(size);
  
  if (!blk)
  {
    printf("malloc(): No free block found!\n");
    return NULL;
  }

  split_block(blk, size);
  blk->free = 0;
  return (void*)((uint8_t*)blk + BLOCK_SIZE);
}

void free(void* ptr)
{
  if (!ptr)
  {
    return;
  }

  block_t* blk = (block_t*)((uint8_t*)ptr - BLOCK_SIZE);
  blk->free = 1;

  /* coalesce */
  block_t* curr = free_list;
  while (curr && curr->next)
  {
    if (curr->free && curr->next->free)
    {
      curr->size += BLOCK_SIZE + curr->next->size;
      curr->next = curr->next->next;
    }
    else
    {
      curr = curr->next;
    }
  }
}

void* calloc(size_t nmemb, size_t size)
{
  size_t total = nmemb * size;
  void* ptr = malloc(total);
  
  if (ptr)
  {
    MEMSET(ptr, 0, total);
  }
  
  return ptr;
}

void* realloc(void* ptr, size_t size)
{
  if (!ptr)
  {
    return malloc(size);
  }
  
  if (!size)
  {
    free(ptr);
    return NULL;
  }

  block_t* blk = (block_t*)((uint8_t*)ptr - BLOCK_SIZE);
  if (blk->size >= size)
  {
    return ptr;
  }

  void* new_ptr = malloc(size);
  
  if (new_ptr)
  {
    memcpy(new_ptr, ptr, blk->size);
    free(ptr);
  }
  return new_ptr;
}