#include <logger.h>
#include <errcode.h>
#include <utils.h>
#include <kmalloc.h>
#include <stddef.h>
#include <mman.h>

extern uint64_t __heap_start;
extern uint64_t __heap_end;

void *(*kmalloc)(size_t size) = simple_alloc;
void (*kfree)(void *ptr) = (void (*)(void *))0x0;

kmalloc_allocator_t mman_kmalloc_pool[KMALLOC_MAX_ALLOCATOR_SIZE + 1];

void *_heap_top = (void *)0x0;

void init_mman_kmalloc()
{
  for (int i = 1; i <= KMALLOC_MAX_ALLOCATOR_SIZE; ++i)
    mman_kmalloc_pool[i] = (kmalloc_allocator_t){
      .left = 0,
      .page_begin = (kmalloc_header_t *)0x0,
    };

  kmalloc = mman_alloc;
  kfree   = mman_free;
}

// simple 8-byte aligned linear allocation
void *simple_alloc(size_t size)
{
  if (!_heap_top) {
    _heap_top = (void *)&__heap_start;
  }

  size = (size_t)ALIGN8(size);
  
  if ((uint64_t)_heap_top + size >= (uint64_t)&__heap_end)
    exit(ERR_NO_MEM);

  void *ret = _heap_top;
  _heap_top = (void *)((uint64_t)_heap_top + size);

  return ret;
}

#define PAGE_HEADER(ptr) ((kmalloc_header_t *)(((uint64_t)(ptr) >> 12) << 12))

void *mman_alloc(size_t unit)
{
  unit = (size_t)ALIGN8(unit);
  if (unit > (KMALLOC_MAX_ALLOCATOR_SIZE << 3))
    return allocate_page(((unit - 1) >> 12) + 1);

  kmalloc_allocator_t *pool = &mman_kmalloc_pool[unit >> 3];
  if (!pool->left) {
    void *page = allocate_page(1), *ptr = page + sizeof(kmalloc_header_t);
    *(kmalloc_header_t *)page = (kmalloc_header_t){
      .unit = unit,
      .left = (PAGE_SIZE - sizeof(kmalloc_header_t)) / unit,
      .begin = page + sizeof(kmalloc_header_t),
      .page_prev = (kmalloc_header_t *)0x0,
      .page_next = (kmalloc_header_t *)0x0,
    };
    *pool = (kmalloc_allocator_t){
      .left = ((kmalloc_header_t *)page)->left,
      .page_begin = (kmalloc_header_t *)page,
    };

    for (; ptr + unit <= page + PAGE_SIZE; ptr += unit)
      *(void **)ptr = ptr + unit;
  }

  void *ret = pool->page_begin->begin;
  pool->page_begin->begin = *(void **)pool->page_begin->begin;
  
  --pool->left;
  --pool->page_begin->left;
  if (!pool->page_begin->left)
    pool->page_begin = *(void **)pool->page_begin->page_next;

  return ret;
}

void mman_free(void *ptr)
{
  if (ptr == ALIGN4K(ptr)) {
    free_page(ptr);
    return;
  }

  size_t unit = PAGE_HEADER(ptr)->unit;
  kmalloc_allocator_t *pool = &mman_kmalloc_pool[unit >> 3];
  if (!PAGE_HEADER(ptr)->left) {
    PAGE_HEADER(ptr)->page_next = pool->page_begin;
    pool->page_begin->page_prev = PAGE_HEADER(ptr);
    pool->page_begin = PAGE_HEADER(ptr);
  }

  *(void **)ptr = pool->page_begin->begin;
  pool->page_begin->begin = ptr;
  
  ++pool->left;
  ++PAGE_HEADER(ptr)->left;
  size_t cap = (PAGE_SIZE - sizeof(kmalloc_allocator_t)) / unit;
  if (PAGE_HEADER(ptr)->left == cap && pool->left >= 2 * cap) {
    if (PAGE_HEADER(ptr)->page_prev)
      PAGE_HEADER(ptr)->page_prev->page_next = PAGE_HEADER(ptr)->page_next;
    if (PAGE_HEADER(ptr)->page_next)
      PAGE_HEADER(ptr)->page_next->page_prev = PAGE_HEADER(ptr)->page_prev;
    free_page(PAGE_HEADER(ptr));
  }
}