#include "arena.h"
#include "base_core.h"
#include "platform.h"


MemoryBlock *memoryblock_allocate(size_t size, size_t commit_size)
{
    size_t aligned_size        = AlignToPageEnd(size);
    size_t aligned_commit_size = AlignToPageEnd(commit_size);

    MemoryBlock *block = memory_reserve(aligned_size);
    Assert(block != NULL);
    memory_commit(block, aligned_commit_size);

    block->base          = (U8 *)block;
    block->size          = aligned_size;
    block->pos           = MemoryBlock_Header_Size;
    block->base_pos      = 0;
    block->prev          = NULL;
    block->committed_pos = aligned_commit_size;

    return block;
}

void arena_init(Arena *arena, const ArenaParams *params)
{
    Assert(params->size > MemoryBlock_Header_Size);
    Assert(params->size < (U32_Max -  platform_page_size())); // keep last page for overflow detection
    Assert(IsPow2(params->commit_size));

    MemoryBlock *block = memoryblock_allocate(params->size, params->commit_size);
    ASAN_UNPOISON_MEMORY_REGION(block, MemoryBlock_Header_Size);

    arena->flags            = params->flags;
    arena->checkpoint_level = 0;
    arena->current          = block;
    arena->block_size       = params->size;
    arena->commit_size      = params->commit_size;

    ASAN_POISON_MEMORY_REGION(&block->base[block->pos], memoryblock_available(block));
}

void arena_free(Arena *arena)
{
    MemoryBlock *block = arena->current;

    while (block != NULL) {
        MemoryBlock *prev = block->prev;
        memory_free(block->base, block->size);
        block = prev;
    }
}

void *arena_push(Arena *arena, U32 size, U32 align)
{
    Assert(size < (U32_Max - platform_page_size()));
    Assert(IsPow2(align));

    MemoryBlock *block = arena->current;
    U32 start_pos = AlignUpPow2(block->pos, align);
    U32 available = block->size - start_pos;

    if (available < size) {
        U32 block_size = arena->block_size; // original user requested size
        U32 new_commit_size = arena->commit_size;
        if (size > block_size - MemoryBlock_Header_Size) {
            block_size = AlignUpPow2(size + MemoryBlock_Header_Size,align);
            new_commit_size = block_size;
        }

        MemoryBlock *new_block  = memoryblock_allocate(block_size, new_commit_size);
        new_block->base_pos     = block->base_pos + block->size;
        new_block->prev         = block;
        arena->current          = new_block;
        block                   = new_block;
        start_pos               = AlignUpPow2(block->pos, align);
    }

    U32 end_pos = start_pos + size;
    Assert(end_pos <= block->size);

    U8 *mem = &block->base[start_pos];
    if (block->committed_pos < end_pos) {
        // Commit more memory
        U32 aligned_end_pos   = AlignToPageEnd(start_pos + size);

        U32 aligned_size      = RoundUpToMultiple(aligned_end_pos - block->committed_pos, arena->commit_size);
        memory_commit(&block->base[block->committed_pos], aligned_size);
        block->committed_pos  += aligned_size;
    }
    block->pos  = end_pos;

    ASAN_UNPOISON_MEMORY_REGION(mem, size);
    return mem;
}

void arena_pop_to(Arena *arena, U32 abs_pos)
{
    MemoryBlock *block = arena->current;
    abs_pos = Max(MemoryBlock_Header_Size, abs_pos);
    while (block) {
        MemoryBlock *prev = block->prev;

        if (block->base_pos >= abs_pos) {
            ASAN_UNPOISON_MEMORY_REGION(block->base, block->size);
            memory_free(block->base, block->size);
            arena->current = prev;
        } else {
            U32 new_pos = abs_pos - block->base_pos;
            new_pos = Max(MemoryBlock_Header_Size, new_pos);
            Assert(new_pos <= block->committed_pos);
            block->pos = new_pos;
            ASAN_POISON_MEMORY_REGION(&block->base[block->pos], memoryblock_available(block));
            break;
        }

        block = prev;
    }
}

ArenaCheckPoint arena_temp_begin(Arena *arena)
{
    ArenaCheckPoint checkpoint = {};
    checkpoint.arena           = arena;
    checkpoint.absolute_pos    = arena_used(arena);
    arena->checkpoint_level++;

    return checkpoint;
}

void arena_temp_end(ArenaCheckPoint *checkpoint)
{
    Arena *arena = checkpoint->arena;
    Assert(arena->checkpoint_level > 0);
    arena_pop_to(arena, checkpoint->absolute_pos);
    arena->checkpoint_level--;
}

void arena_reset(Arena *arena)
{
    arena_pop_to(arena, 0);
}