#include "buddy.h"
#include "arena.h"

/* TODO: Elminiate meta arena. Just store levels in first 4096 bytes of buddy arena, and then zone header data at the beginning of each allocation. */

/* ========================== *
 * Ctx
 * ========================== */

struct buddy_ctx *buddy_ctx_alloc(u64 reserve)
{
    /* TODO: Determine meta reserve dynamically */
    struct arena meta_arena = arena_alloc(GIGABYTE(64));
    struct buddy_ctx *ctx = arena_push_zero(&meta_arena, struct buddy_ctx);
    ctx->meta_arena = meta_arena;
    ctx->data_arena = arena_alloc(reserve);

    /* TODO: Minimum block size */
    ctx->levels = arena_push_array_zero(&ctx->meta_arena, struct buddy_level, 64);
    for (u64 i = 0; i < 64; ++i) {
        struct buddy_level *level = &ctx->levels[i];
        level->ctx = ctx;
        level->tier = i;
        level->size = (u64)1 << i;
    }

    return ctx;
}

void buddy_ctx_release(struct buddy_ctx *ctx)
{
    arena_release(&ctx->data_arena);
    arena_release(&ctx->meta_arena);
}

/* ========================== *
 * Block
 * ========================== */

INTERNAL struct buddy_block *buddy_block_alloc_internal(struct buddy_ctx *ctx)
{
    struct buddy_block *block;
    if (ctx->first_free_block) {
        block = ctx->first_free_block;
        ctx->first_free_block = block->next;
    } else {
        block = arena_push(&ctx->meta_arena, struct buddy_block);
    }
    MEMZERO_STRUCT(block);
    return block;
}

INTERNAL void buddy_block_release_internal(struct buddy_ctx *ctx, struct buddy_level *level, struct buddy_block *block)
{
    /* Remove from unused list */
    {
        struct buddy_block *prev = block->prev;
        struct buddy_block *next = block->next;
        if (prev) {
            prev->next = next;
        } else {
            level->first_unused_block = next;
        }
        if (next) {
            next->prev = prev;
        }
    }

    block->next = ctx->first_free_block;
    ctx->first_free_block = block;
}

INTERNAL struct buddy_block *buddy_block_get_unused(struct buddy_ctx *ctx, struct buddy_level *level)
{
    struct buddy_block *block = NULL;

    /* TODO: Tier oob check */
    if (level->first_unused_block) {
        block = level->first_unused_block;
        level->first_unused_block = block->next;
        if (level->first_unused_block) {
            level->first_unused_block->prev = NULL;
        }
        block->used = true;
        block->next = NULL;
    } else {
        if (level->backed) {
            struct buddy_level *parent_level = &ctx->levels[level->tier + 1];
            struct buddy_block *parent_block = buddy_block_get_unused(ctx, parent_level);

            /* Create left (used) block from parent block */
            struct buddy_block *left = buddy_block_alloc_internal(ctx);
            left->used = true;
            left->level = level;
            left->parent = parent_block;
            left->memory = parent_block->memory;

            /* Create right (unused) block from parent block */
            struct buddy_block *right = buddy_block_alloc_internal(ctx);
            right->used = false;
            right->level = level;
            right->parent = parent_block;
            right->memory = left->memory + level->size;
            if (level->first_unused_block) {
                right->next = level->first_unused_block;
                level->first_unused_block->prev = right;
            }
            level->first_unused_block = right;

            left->sibling = right;
            right->sibling = left;
            block = left;
        } else {
            struct arena *arena = &ctx->data_arena;

            /* Grow arena */
            i64 level_commit_diff = (level->size * 2) - arena->pos;
            if (level_commit_diff > 0) {
                arena_push_array(arena, u8, level_commit_diff);
                ASSERT(arena->pos == (level->size * 2));
            }

            /* Create left (used) block from existing child block memory */
            struct buddy_block *left = buddy_block_alloc_internal(ctx);
            left->used = true;
            left->level = level;
            left->memory = arena->base;

            /* Create right (unused) block from new arena memory */
            struct buddy_block *right = buddy_block_alloc_internal(ctx);
            right->used = false;
            right->level = level;
            right->memory = left->memory + level->size;

            left->sibling = right;
            right->sibling = left;
            block = left;

            level->backed = true;
        }
    }

    return block;
}

INTERNAL void buddy_block_mark_unused(struct buddy_block *block)
{
    block->used = false;
    struct buddy_level *level = block->level;
    struct buddy_block *parent = block->parent;
    struct buddy_block *sibling = block->sibling;
    if (!sibling->used && parent != NULL) {
        /* Merge siblings */
        struct buddy_ctx *ctx = level->ctx;
        buddy_block_release_internal(ctx, level, block);
        buddy_block_release_internal(ctx, level, sibling);
        buddy_block_mark_unused(parent);
    } else {
        if (level->first_unused_block) {
            block->next = level->first_unused_block;
            level->first_unused_block->prev = block;
        }
        level->first_unused_block = block;
    }
}

struct buddy_block *buddy_alloc(struct buddy_ctx *ctx, u64 size)
{
    if (size > 0x00FFFFFFFFFFFFFFULL) {
        /* TODO: Error */
        ASSERT(false);
    }

    /* TODO: Minimum block size */

    u64 desired_block_size = 1;
    u64 desired_level_tier = 0;
    while (desired_block_size < size && desired_level_tier < 64) {
        desired_block_size <<= 1;
        ++desired_level_tier;
    }

    struct buddy_block *block = buddy_block_get_unused(ctx, &ctx->levels[desired_level_tier]);

    return block;
}

void buddy_release(struct buddy_block *block)
{
    buddy_block_mark_unused(block);
}