#include <stdio.h>
#include <string.h>
#include <assert.h>
#include "buddy.h"

static uint8_t MAGIC[MAGIC_COOKIE_BYTES];

struct free_list {
    uint8_t magic[MAGIC_COOKIE_BYTES];
    struct free_list *next;
    struct free_list **pprev;
};

static struct free_list *(AVAIL[ADDRESS_BITS]);
static size_t ROOT_LOGSIZE;
static void *BASE;

// log2 of size rounded up (ceil) or down (!ceil) to the nearest power of 2
static size_t size2log(size_t size, int ceil) {
    int bitcnt = 0;
    for (int i = 0; size; i++) {
        bitcnt += size & 1;
        size >>= 1;
        if (size) continue;
        if (ceil) return (bitcnt > 1) ? (i + 1) : i;
        return i;
    }
    return ceil ? 1 : 0;
}

void init_buddy(uint8_t *base, size_t size,
                uint8_t magic[MAGIC_COOKIE_BYTES]) {
    BASE = base;

    size_t logsize = size2log(size, 0);
    ROOT_LOGSIZE = logsize;

    AVAIL[logsize] = (struct free_list *)base;
    AVAIL[logsize]->next = 0;
    AVAIL[logsize]->pprev = &(AVAIL[logsize]);
    memcpy(MAGIC, magic, sizeof(MAGIC));
}

static struct free_list *buddy_of(struct free_list *block, size_t logsize) {
    size_t virt_block = (size_t)block - (size_t)BASE;
    size_t virt_buddy = virt_block ^ (1 << logsize);
    return (void*)((uint8_t*)BASE + virt_buddy);
}

static int isfree(struct free_list *block) {
    return !memcmp(block->magic, MAGIC, sizeof(MAGIC));
}

static void makefree(struct free_list *block) {
    memcpy(block->magic, MAGIC, sizeof(MAGIC));
}

static struct free_list *pop(struct free_list *block) {
    *(block->pprev) = block->next;
    if (block->next) block->next->pprev = block->pprev;
    return block;
}

static void push(struct free_list *block, size_t logsize) {
    block->next = AVAIL[logsize];
    if (block->next) block->next->pprev = &(block->next);
    AVAIL[logsize] = block;
    block->pprev = &(AVAIL[logsize]);
}

static void *_allocate(size_t logsize) {
    if (logsize > ROOT_LOGSIZE) return 0;
    if (AVAIL[logsize]) {
        struct free_list *block = pop(AVAIL[logsize]);
        memset(block, 0, sizeof(struct free_list));
        return block;
    }

    struct free_list *parent = _allocate(logsize + 1);
    struct free_list *buddy = buddy_of(parent, logsize);
    // split @parent in half and place the buddy on the avail list.
    memcpy(buddy->magic, MAGIC, sizeof(MAGIC));
    push(buddy, logsize);
    return parent;
}
void *allocate(size_t size) {
    size = (size < sizeof(struct free_list)) ? sizeof(struct free_list) : size;
    return _allocate(size2log(size, 1));
}

static void _liberate(struct free_list *block, size_t logsize) {
    push(block, logsize);
    if (logsize == ROOT_LOGSIZE) return;

    struct free_list *buddy = buddy_of(block, logsize);
    if (!isfree(buddy)) return;

    // coalesce up!
    struct free_list *smaller = (buddy < block) ? buddy : block;
    struct free_list *bigger = (buddy < block) ? block : buddy;
    pop(bigger);
    memset(bigger, 0, sizeof(*bigger));

    pop(smaller);

    _liberate(smaller, logsize + 1);
}

void liberate(void *base, size_t size) {
    struct free_list *block = base;
    memset(block, 0, sizeof(*block));
    makefree(block);

    _liberate(block, size2log(size, 1));
}

void debug_buddy(void) {
    for (size_t i = 0; i <= ROOT_LOGSIZE; i++) {
        printf("Free blocks at size 2^%lu = %lu:\n", i, 1ul << i);
        for (struct free_list *block = AVAIL[i]; block; block = block->next) {
            printf("\t%p\n", block);
        }
    }
}

///////// "advanced features"
static struct free_list *rhs_child_of(struct free_list *block, size_t logsize) {
    size_t virt_block = (size_t)block - (size_t)BASE;
    size_t virt_child = virt_block | (1 << (logsize - 1));
    return (void*)((uint8_t*)BASE + virt_buddy);
}

// NOTE: this method is perhaps more complicated than it needs to be because we
// take great pains to avoid writing to the region that is being preallocated
// (e.g., in case it is device MMIO).
int preallocate(void *start, size_t size) {
    assert(size);
    void *end = (void*)((uint8_t*)start + size);

    // (1) check if BASE itself is a free node
    struct free_list *parent = (void*)BASE;
    size_t parent_logsize = ROOT_LOGSIZE;
    while (!isfree(parent)) {
        // pick the proper child and recurse
        struct free_list *rhs_child = rhs_child_of(parent, parent_logsize);
        if (rhs_child <= start) parent = rhs_child;
        if (!parent_logsize--) return 0;
    }
    // parent is free!
    void *parent_end = (void*)((uint8_t*)parent + (1 << parent_logsize));
    // if the size has no chance of fitting, return 0.
    if (parent_end < end) return 0;
    // otherwise, split this parent
    pop(parent, parent_logsize);
    while (1) {
        // check whether the child could fit it.
        struct free_list *rhs_child = rhs_child_of(parent, parent_logsize);
        struct free_list *next_child = (rhs_child <= start) parent : rhs_child;
        void *child_end = (void*)((uint8_t*)next_child + (1 << (parent_logsize - 1)));
        if (child_end < end) {
            // the child *cannot* fit it, so allocate this parent.
            return 1;
        } else {
            // the child *can* fit it, so split this parent into two children.
            struct free_list *lhs_child = parent;
            if (rhs_child <= start) {
                // the lhs child will be free
                makefree(lhs_child);
                push(lhs_child, parent_logsize - 1);
                parent = rhs_child;
            } else {
                // the rhs child will be free
                makefree(rhs_child);
                push(rhs_child, parent_logsize - 1);
                parent = lhs_child;
            }
        }
        parent_logsize--;
    }
    return 0;
}

void *reallocate(void *old, size_t old_size, size_t new_size) {
    // if new_size < old_size, repeatedly split ourself. TODO: does this mess
    // with fragmentation?

    // otherwise, basically identical to preallocate, except what we want to
    // allocate is roughly @old + old_size (with proper rounding to
    // boundaries).
}