#include "malloc.h"
#include <assert.h>
#include <stdbool.h>
#include <unistd.h>
#include <errno.h>

/**Allocate memory for a new bucket in the heap.
 * Automatically aligns new memory to BUCKET_ALIGN.
 * Instantiates a new free object spanning the entire bucket.
 * Then adds the new bucket to the heap cache.
 * If sbrk encounters an error, pass through errno and return NULL.
 * @param cache to allocate for.
 * @param objw width of the object the bucket is assumed to hold.
 * @return bucket created by this function.*/
bucket_t*
s_bucket_new(heap_cache_t *cache, size_t objw)
{
    DBGPRINT("malloc: ask for new bucket for cache %p with objw %i\n", cache, objw);
    size_t bucketw = BUCKET_WIDTH(objw);
    /* Assuming sbrk allocates memory for after the program break, make sure
     * we allocate enough for the requested bucket width and any adjustments
     * to ensure page alignment.*/
    uintptr_t pbrk = (uintptr_t)sbrk(0);
    uintptr_t brko = pbrk % BUCKET_ALIGN;
    uintptr_t nbrk = (uintptr_t)sbrk(bucketw + brko);
    if(errno != 0) return NULL; /* Fail on errno. */

    bucket_t *bucket = (bucket_t*)(nbrk + brko);
    DBGPRINT("malloc: got new bucket at %p\n", bucket);

    *bucket = (bucket_t){
        .next = NULL,
        .limit = (uintptr_t)bucket + bucketw,
    };
    /* Create a new object spanning the entire bucket. This new object will
     * be the first free.*/
    bucket_obj_t *obj0 = (bucket_obj_t*)bucket->objs;
    *obj0 = (bucket_obj_t) {
        .checksum = OBJ_CHECKSUM,
        .size_taken = bucketw - sizeof(bucket_t) - sizeof(bucket_obj_t)
    };
    bucket->firstfree = obj0;

    /* If the cache is empty, this will be the first and last bucket.*/
    if(cache->head == NULL) {
        cache->head = cache->tail = bucket;
        return bucket;
    }
    /* Add the bucket to the heap cache bucket list.*/
    cache->tail->next = bucket;
    cache->tail = bucket;
    return bucket;
}

/**Given an existing object and a size, split the object so that the original
 * has size w.
 * If the object is smaller than or is already the correct size, do nothing.
 * Space after the split is allocated to a new free object after the resized
 * original.
 * If the original == bucket->firstfree, the new object is assigned to 
 * that variable.
 * If the heap is corrupted, assert fail.
 * @param bucket object's parent bucket.
 * @param parent object to split.
 * @param w requested size for original object. */
void
s_obj_take_split(bucket_t *bucket, bucket_obj_t *parent, size_t w)
{
    DBGPRINT("malloc: requesting to split object %p (last %p) for size %i (was %i)\n",
           bucket, parent, w, parent->size_taken);

    /* Assert fail on corrupted heap.*/
    assert(OBJ_VALID(bucket));
    assert(OBJ_VALID(parent));
    assert(w >= OBJ_SIZE_MIN);

    size_t objw = OBJ_WIDTH(parent);
    if(w == objw) { /* Obj is already correct size, take and ignore.*/
        if(bucket->firstfree == parent)
            bucket->firstfree = OBJ_AFTER(parent);
        /* Ensure that firstfree does not exceed bucket limits.*/
        if((uintptr_t)bucket->firstfree > bucket->limit)
            bucket->firstfree = NULL;
        DBGPRINT("malloc: Object is already right size, ignoring\n");

        /* Mark object as taken.*/
        parent->size_taken = w | 1;
        return;
    }
    /* New object created by split will be the original width minus the new
     * width and metadata size.
     *
     * from:
     * [meta| --------objw-------- ]
     *
     * to:
     * [meta| -w- ][meta| -afterw- ]*/
    size_t afterw = objw - sizeof(bucket_obj_t) - w;

    parent->size_taken = w | 1;
    bucket_obj_t *obj = OBJ_AFTER(parent);
    DBGPRINT("malloc: New object created inplace at %p with size %i\n",
            obj, afterw);

    /* In-place allocate after obj.*/
    *obj = (bucket_obj_t){
        .checksum = OBJ_CHECKSUM,
        .size_taken = afterw,
    };
    if(bucket->firstfree == parent)
        bucket->firstfree = obj;
}

/* Look for a free object in the given bucket that closest matches size w.
 * Uses bucket->firstfree as a starting point to save checks against taken
 * objects.
 * If the heap is corrupted, assert fail.
 * If there are no good objects in this bucket, return NULL.
 * @param bucket to search in.
 * @param w to search for.
 * @return best-fitting object or NULL if none are found.
 * */
bucket_obj_t*
s_obj_bestfit(bucket_t *bucket, size_t w)
{
    assert(bucket != NULL);
    assert(OBJ_VALID(bucket));

    bucket_obj_t *best = NULL;
    bucket_obj_t *obj = bucket->firstfree;
    if(obj == NULL) return NULL;
    for(; (uintptr_t)obj < bucket->limit; obj = OBJ_AFTER(obj))
    {
        assert(OBJ_VALID(obj));
        if(OBJ_TAKEN(obj)) continue;
        size_t objw = OBJ_WIDTH(obj);
        if(objw < w) continue;
        if(best == NULL) best = obj;

        size_t bestw = OBJ_WIDTH(best);
        if(objw < bestw) best = obj;
    }
    
    DBGPRINT("malloc: Bucket %p returns best fit at %p for size %i ",
            bucket, best, w);
    if(best != NULL) { DBGPRINT("(size %i)\n", best->size_taken); }
    else { DBGPRINT("(is NULL)\n"); }

    return best;
}

/**Searching through the entire heap, or until an object is found, find a
 * free object that best fits the requested width and place the resulting
 * bucket in the passed pointer.
 * If the heap is corrupted, assert fail.
 * If no object is found in any of the buckets, return NULL.
 * @param cache heap to search through.
 * @param w object width to search for.
 * @param bucket pointer to bucket pointer owning the returned bucket.
 *        If no object is found, variable is untouched.
 * @return object that fits the given width.*/
bucket_obj_t*
s_bucket_bestfit(heap_cache_t *cache, size_t w, bucket_t **bucket)
{
    *bucket = cache->head;
    for(; *bucket != NULL; *bucket = (*bucket)->next)
    {
        bucket_obj_t *obj = s_obj_bestfit(*bucket, w);
        /* Object is already found, don't waste cycles searching through
         * the entire heap.*/
        if(obj != NULL) return obj;
    }
    return NULL;
}

void*
__malloc_impl(heap_cache_t *cache, size_t w)
{
    assert(cache != NULL);
    /* w == 0 is ub, just return NULL.*/
    if(w == 0) return NULL;
    /* Ensure that w is valid.*/
    w = REALW(w);

    DBGPRINT("malloc: Asking to malloc for size %i\n", w);

    bucket_t *bucket = NULL;
    bucket_obj_t *obj = s_bucket_bestfit(cache, w, &bucket);
    if(obj == NULL) {
        /* No bucket in the heap has an appropriate space.
         * Allocate a new bucket and take the first object.
         * If the bucket allocation fails, return -1 and pass errno through.*/
        bucket = s_bucket_new(cache, w);
        if(bucket == NULL) return (void*)(-1);
        obj = (bucket_obj_t*)bucket->objs;
    }
    /* Split the bucket to make the object bestfit.*/
    s_obj_take_split(bucket, obj, w);
    return obj->data;
}