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


static MemoryBlock *memoryblock_allocate(U32 size, U32 commit_size, U32 page_size)
{
    Assert(IsAlignedTo(size, page_size));
    Assert(IsAlignedTo(commit_size, page_size));

    MemoryBlock *block = memory_reserve(size);
    Assert(block != NULL);
    memory_commit(block, commit_size);

    ASAN_POISON_MEMORY_REGION(block, size);
    ASAN_UNPOISON_MEMORY_REGION(block, MemoryBlock_Header_Size);

    block->base          = (U8 *)block;
    block->size          = size;
    block->first_free    = MemoryBlock_Header_Size;
    block->arena_offset  = 0;
    block->prev          = NULL;
    block->committed_end = commit_size;
    return block;
}

void arena_init(Arena *arena, const ArenaOptions *opts)
{
    U32 page_size = platform_page_size();
    Assert(opts->size > MemoryBlock_Header_Size);
    Assert(opts->size < (U32_Max -  page_size)); // Reserving last piece so our math doesn't overflow
    Assert(opts->commit_size > 0);

    U32 aligned_commit_size = AlignUp(opts->commit_size, page_size);
    U32 aligned_size        = AlignUp(opts->size, page_size);
    Assert(aligned_commit_size <= aligned_size);

    MemoryBlock *block  = memoryblock_allocate(aligned_size, aligned_commit_size, page_size);

    arena->flags            = opts->flags;
    arena->temp_level = 0;
    arena->current          = block;
    arena->block_size       = aligned_size;
    arena->commit_size      = aligned_commit_size;
    arena->page_size        = page_size;
}

void arena_free(Arena *arena)
{
    MemoryBlock *block = arena->current;
    Assert(arena->temp_level == 0);
    while (block != NULL) {
        MemoryBlock *prev = block->prev;
        memory_free(block->base, block->size);
        block = prev;
    }
    arena->current = NULL;
}

void *arena_push(Arena *arena, U32 size, U32 align)
{
    Assert(IsPow2(align));


    MemoryBlock *block = arena->current;
    U32 start_pos = AlignUp(block->first_free, align);  // might align past block->size, handled in below check
    U32 available = block->size - start_pos;

    if (start_pos > block->size || available < size) {
        U32 pad = AlignUp(MemoryBlock_Header_Size, align);
        Assert(size <= (U32_Max - arena->page_size) - pad);
        U32 needed = pad + size;
        U32 new_block_size = arena->block_size;
        U32 new_commit_size = arena->commit_size;
        if (needed > new_block_size) {
            new_block_size = AlignUp(needed, arena->page_size);
            new_commit_size = new_block_size;
        }

        MemoryBlock *new_block = memoryblock_allocate(new_block_size, new_commit_size, arena->page_size);
        // arena_offset tracks the virtual absolute position within an arena
        // This is not a contiguous space, we use it to free blocks when popping the arena
        new_block->arena_offset = block->arena_offset + block->size;
        new_block->prev         = block;
        arena->current          = new_block;
        block                   = new_block;
        start_pos               = AlignUp(block->first_free, align);
    }

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

    U8 *mem = &block->base[start_pos];
    if (block->committed_end < end_pos) {
        // Commit more memory
        U32 aligned_end_pos   = AlignUp(start_pos + size, arena->page_size);
        U32 aligned_size      = RoundUpToMultiple(aligned_end_pos - block->committed_end, arena->commit_size);
        aligned_size = Min(aligned_size, block->size - block->committed_end);

        memory_commit(&block->base[block->committed_end], aligned_size);
        block->committed_end  += aligned_size;
        Assert(block->committed_end <= block->size);
    }
    block->first_free  = end_pos;

    ASAN_UNPOISON_MEMORY_REGION(mem, size);
    return mem;
}

void arena_reset_to(Arena *arena, U64 total_offset)
{
    MemoryBlock *block = arena->current;
    total_offset = Max(MemoryBlock_Header_Size, total_offset);
    while (block) {
        MemoryBlock *prev = block->prev;

        if (block->arena_offset >= total_offset) {
            ASAN_UNPOISON_MEMORY_REGION(block->base, block->size);
            memory_free(block->base, block->size);
            arena->current = prev;
        } else {
            U64 new_pos = total_offset - block->arena_offset;
            new_pos = Max(MemoryBlock_Header_Size, new_pos);
            // Disallow arbitrary pops to positions that previously were not commited.
            Assert(new_pos <= block->committed_end);
            block->first_free = new_pos;
            ASAN_POISON_MEMORY_REGION(&block->base[block->first_free], memoryblock_available(block));
            break;
        }

        block = prev;
    }
}

ATemp arena_temp_begin(Arena *arena)
{
    ATemp checkpoint = {};
    checkpoint.arena           = arena;
    checkpoint.total_offset    = arena_used(arena);
    arena->temp_level++;

    return checkpoint;
}

void arena_temp_end(ATemp *checkpoint)
{
    Arena *arena = checkpoint->arena;
    Assert(arena->temp_level > 0);
    arena_reset_to(arena, checkpoint->total_offset);
    arena->temp_level--;
}

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