#include "heap.h"
#include <jove/memory.h>
#include <jove/jove.h>
#include <stdlib.h>
#include <string.h>

static heap_bin_t *freelist_tail;

uintptr_t __heap_start;
uintptr_t __heap_end;

#define HEAP_BIN_SIZE(bin) (bin->size_taken & (~1ULL))
#define HEAP_BIN_TAKEN(bin) (bin->size_taken & 1)

#define HEAP_BIN_AFTER(bin) ((heap_bin_t*)((uintptr_t)bin->data + HEAP_BIN_SIZE(bin) + sizeof(uintptr_t)))
#define HEAP_BIN_BEFORE(bin) ((heap_bin_t*)*((uintptr_t*)bin - 1))

static void
heap_grow(uintptr_t to)
{
    jove_mem_ensure_range(__heap_end, to);
    __heap_end = to;
}

static heap_bin_t*
heap_newbin(uintptr_t at, size_t size, heap_bin_t *next)
{
    uintptr_t headat = (at + sizeof(heap_bin_t) + size);
    uintptr_t endat = headat + sizeof(uintptr_t);

    if(endat > __heap_end) heap_grow(endat);

    heap_bin_t *newbin = (heap_bin_t*)at;
    newbin->size_taken = size;
    newbin->next = next;
    
    *((uintptr_t*)headat) = at;
    return newbin;
}

static heap_bin_t*
heap_find_bestfit(size_t rsize)
{
    heap_bin_t *best = freelist_tail;
    if(best == NULL) return NULL;
    size_t bestsize = best->size_taken;

    for(heap_bin_t *bin = best->next; bin; bin = bin->next)
    {
        if(bin->size_taken < rsize) continue;
        if(bin->size_taken < bestsize) {
            best = bin;
            bestsize = bin->size_taken;
        }
    }
    return best;
}

static void
heap_split(heap_bin_t *bin, size_t splitsize)
{
    size_t size_diff = bin->size_taken - splitsize;
    if(size_diff < LIBC_HEAP_MINSPLIT) return;

    size_t newsize = size_diff - sizeof(heap_bin_t) - sizeof(uintptr_t);
    uintptr_t newat = ((uintptr_t)bin->data) + splitsize + sizeof(uintptr_t);
    jove_kprintf("Split %p to %p [%x]\n", bin, newat, newsize);

    bin->size_taken = splitsize;
    freelist_tail = heap_newbin(newat, newsize, freelist_tail);
}

static void*
heap_alloc(size_t size)
{
    if(size & 1) size++;
    if(size < LIBC_HEAP_MINALLOC) size = LIBC_HEAP_MINALLOC;

    heap_bin_t *bestfit = heap_find_bestfit(size);
    if(bestfit == NULL) {
        bestfit = heap_newbin(__heap_end, size, NULL);
    }else{
        if(bestfit == freelist_tail) {
            freelist_tail = bestfit->next;
        }
        size_t bestsize = bestfit->size_taken;

        if(bestsize < size) {
            jove_kprintf("Growing bin %p [%x] to %x\n", bestfit, bestsize, size);
            heap_newbin((uintptr_t)bestfit, size, NULL);
        }else if(bestsize > size) {
             heap_split(bestfit, size);
        }
    }
    bestfit->size_taken |= 1;
    return bestfit->data;
}

static void 
heap_free(void *ptr)
{
    uintptr_t binptr = (uintptr_t)ptr - sizeof(heap_bin_t);
    if(binptr > __heap_end || binptr < __heap_start) return;

    heap_bin_t *bin = (heap_bin_t*)binptr;
    bin->size_taken &= ~1ULL;
    bin->next = freelist_tail;
    freelist_tail = bin;

start_merge_loop: {
    heap_bin_t *prevbin = HEAP_BIN_BEFORE(bin);
    heap_bin_t *nextbin = HEAP_BIN_AFTER(bin);

    if((uintptr_t)prevbin >= __heap_start) {
        if(!HEAP_BIN_TAKEN(prevbin)) {
            jove_kprintf("Merge back bin %p into %p\n", bin, prevbin);
            freelist_tail = prevbin;
            prevbin->size_taken += bin->size_taken + sizeof(heap_bin_t) + sizeof(uintptr_t);
            bin = prevbin;
            goto start_merge_loop;
        }
    }
    if((uintptr_t)nextbin < __heap_end) {
        if(!HEAP_BIN_TAKEN(nextbin)) {
            jove_kprintf("Merge forward bin %p into %p\n", bin, nextbin);
            if(bin->next == nextbin) bin->next = nextbin->next;
            bin->size_taken += nextbin->size_taken + sizeof(heap_bin_t) + sizeof(uintptr_t);
            goto start_merge_loop;
        }
    }
                  }
}

void
__libc_heap_init(uintptr_t program_end)
{
    __heap_start = __heap_end = program_end;
    freelist_tail = heap_newbin(__heap_start, LIBC_HEAP_INIT_SIZE, NULL);
}