#include <buddy.h>
#include <os_macros.h>
#include <string.h>
#include <os_compiler.h>


#define IS_POWER_OF_2(n) ((n) != 0 && (((n) & ((n) - 1)) == 0))

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

OS_STATIC_FORCE_INLINE
buddy_block_t* buddy_block_next(buddy_block_t* block){
    return (buddy_block_t*)((uint8_t*)block + block->size);
}


static
buddy_block_t *buddy_block_split(buddy_block_t *block, size_t size) {
    if (block != NULL && size != 0) {
        // Recursive split
        while (size < block->size) {
            size_t sz = block->size >> 1;
            block->size = sz;
            block = buddy_block_next(block);
            block->size = sz;
            block->is_free = OS_TRUE;
        }

        if (size <= block->size) {
            return block;
        }
    }

    // Block cannot fit the requested allocation size
    return NULL;
}

static
buddy_block_t *buddy_block_find_best(buddy_block_t *head, buddy_block_t *tail, size_t size) {
    // Assumes size != 0

    buddy_block_t *best_block = NULL;
    buddy_block_t *block = head;                    // Left Buddy
    buddy_block_t *buddy = buddy_block_next(block); // Right Buddy

    // The entire memory section between head and tail is free,
    // just call 'buddy_block_split' to get the allocation
    if (buddy == tail && block->is_free) {
        return buddy_block_split(block, size);
    }

    // Find the block which is the 'best_block' to requested allocation sized
    while (block < tail && buddy < tail) { // make sure the buddies are within the range
        // If both buddies are free, coalesce them together
        // NOTE: this is an optimization to reduce fragmentation
        //       this could be completely ignored
        if (block->is_free && buddy->is_free && block->size == buddy->size) {
            block->size <<= 1;
            if (size <= block->size && (best_block == NULL || block->size <= best_block->size)) {
                best_block = block;
            }

            block = buddy_block_next(buddy);
            if (block < tail) {
                // Delay the buddy block for the next iteration
                buddy = buddy_block_next(block);
            }
            continue;
        }


        if (block->is_free && size <= block->size &&
            (best_block == NULL || block->size <= best_block->size)) {
            best_block = block;
        }

        if (buddy->is_free && size <= buddy->size &&
            (best_block == NULL || buddy->size < best_block->size)) {
            // If each buddy are the same size, then it makes more sense
            // to pick the buddy as it "bounces around" less
            best_block = buddy;
        }

        if (block->size <= buddy->size) {
            block = buddy_block_next(buddy);
            if (block < tail) {
                // Delay the buddy block for the next iteration
                buddy = buddy_block_next(block);
            }
        } else {
            // Buddy was split into smaller blocks
            block = buddy;
            buddy = buddy_block_next(buddy);
        }
    }

    if (best_block != NULL) {
        // This will handle the case if the 'best_block' is also the perfect fit
        return buddy_block_split(best_block, size);
    }

    // Maybe out of memory
    return NULL;
}

OS_STATIC_FORCE_INLINE
os_size_t align_forward_size(os_size_t ptr, os_size_t align) {
    os_size_t a, p, modulo;

    OS_ASSERT(IS_POWER_OF_2((os_uintptr_t)align));

    a = align;
    p = ptr;
    modulo = p & (a-1);
    if (modulo != 0) {
        p += a - modulo;
    }
    return p;
}

OS_STATIC_FORCE_INLINE
os_size_t buddy_block_size_required(buddy_t *b, size_t size) {
    os_size_t actual_size = b->alignment;

    size += sizeof(buddy_block_t);
    size = align_forward_size(size, b->alignment);

    while (size > actual_size) {
        actual_size <<= 1;
    }

    return actual_size;
}

static
void buddy_block_coalescence(buddy_block_t *head, buddy_block_t *tail) {
    for (;;) {
        // Keep looping until there are no more buddies to coalesce

        buddy_block_t *block = head;
        buddy_block_t *buddy = buddy_block_next(block);

#if 0
        // 原版有bug
        bool no_coalescence = OS_TRUE;
        while (block < tail && buddy < tail) { // make sure the buddies are within the range
            if (block->is_free && buddy->is_free && block->size == buddy->size) {
                // Coalesce buddies into one
                block->size <<= 1;
                block = buddy_block_next(block);
                if (block < tail) {
                    buddy = buddy_block_next(block);
                    no_coalescence = C_FALSE;
                }
            } else if (block->size < buddy->size) {
                // The buddy block is split into smaller blocks
                block = buddy;
                buddy = buddy_block_next(buddy);
            } else {
                block = buddy_block_next(buddy);
                if (block < tail) {
                    // Leave the buddy block for the next iteration
                    buddy = buddy_block_next(block);
                }
            }
        }
#endif

        bool no_coalescence = OS_TRUE;
        while (block < tail && buddy < tail) { // make sure the buddies are within the range
            if (block->is_free && buddy->is_free && block->size == buddy->size) {
                // Coalesce buddies into one
                block->size <<= 1;
//                block = buddy_block_next(buddy);
                if (block < tail) {
                    buddy = buddy_block_next(block);
                    no_coalescence = OS_FALSE;
                }
            } else if (block->size <= buddy->size) {
                block = buddy_block_next(buddy);
                if (block < tail) {
                    // Leave the buddy block for the next iteration
                    buddy = buddy_block_next(block);
                }
            } else {
                // The buddy block is split into smaller blocks
                block = buddy;
                buddy = buddy_block_next(buddy);
            }
        }

        if (no_coalescence) {
            return;
        }
    }
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

void buddy_init(buddy_t* self, void* block, os_size_t block_size, os_size_t alignment){
    OS_ASSERT(block != NULL);
    OS_ASSERT(IS_POWER_OF_2(block_size) && "size is not a power-of-two");
    OS_ASSERT(IS_POWER_OF_2(alignment) && "alignment is not a power-of-two");

    // The minimum alignment depends on the size of the `budy_block_t` header
    alignment = OS_MAX(alignment, sizeof(buddy_block_t));
//    if (alignment < sizeof(buddy_block_t)) {
//        alignment = sizeof(buddy_block_t);
//    }
    OS_ASSERT((uintptr_t)block % alignment == 0 && "block is not aligned to minimum alignment");

    self->head          = (buddy_block_t *)block;
    self->head->size    = block_size;
    self->head->is_free = OS_TRUE;

    // The tail here is a sentinel value and not a true block
    self->tail = buddy_block_next(self->head);

    self->alignment = alignment;
}

void buddy_destroy(buddy_t* self){
    self->head = 0;
    self->tail = 0;
    self->alignment = 0;
}

void* buddy_alloc(buddy_t* b, os_size_t size){
    if (size != 0) {
        size_t actual_size = buddy_block_size_required(b, size);

        buddy_block_t *found = buddy_block_find_best(b->head, b->tail, actual_size);
        if (found == NULL) {
            // Try to coalesce all the free buddy blocks and then search again
            buddy_block_coalescence(b->head, b->tail);
            found = buddy_block_find_best(b->head, b->tail, actual_size);
        }

        if (found != NULL) {
            found->is_free = OS_FALSE;
            return (void *)((uint8_t *)found + b->alignment);
        }

        // Out of memory (possibly due to too much internal fragmentation)
    }

    return NULL;
}

void* buddy_calloc(buddy_t* b, os_size_t count, os_size_t size){
    void* mem = buddy_alloc(b, size * count);
    if(mem){
        memset(mem, 0, size);
    }
    return mem;
}

void* buddy_realloc(buddy_t* self, void* ptr, os_size_t size){
    void* mem = buddy_alloc(self, size);
    if(!mem){
        return NULL;
    }
    buddy_block_t *block =(buddy_block_t *)((uint8_t *)ptr - self->alignment);
    size = OS_MIN(size, block->size);
    memcpy(mem, ptr, size);
    buddy_free(self, ptr);
    return mem;
}

void buddy_free(buddy_t* b, void* data){
    if (data != NULL) {
        buddy_block_t *block;

        OS_ASSERT(b->head <= (buddy_block_t*)data);
        OS_ASSERT((buddy_block_t*)data < b->tail);

        block = (buddy_block_t *)((uint8_t *)data - b->alignment);
        block->is_free = OS_TRUE;

        // NOTE: Coalescence could be done now but it is optional
//         buddy_block_coalescence(b->head, b->tail);
    }
}