#include "arena.h"
#include "sys.h"
#include "memory.h"
#include "string.h"

/* Arbitrary block size */
#define ARENA_BLOCK_SIZE 4096

/* NOTE: Application will exit if arena fails to reserve or commit initial
 * memory. */
struct arena arena_alloc(u64 reserve)
{
    __prof;
    struct arena arena = { 0 };

    /* Round up to nearest block size */
    u64 block_remainder = reserve % ARENA_BLOCK_SIZE;
    if (block_remainder > 0) {
        reserve += ARENA_BLOCK_SIZE - block_remainder;
    }

    arena.base = sys_memory_reserve(reserve);
    if (!arena.base) {
        /* Hard fail on memory reserve failure for now */
        sys_panic(STR("Failed to reserve memory"));
    }
    arena.reserved = reserve;

    /* Commit one block to start with */
    arena.base = sys_memory_commit(arena.base, ARENA_BLOCK_SIZE);
    __profalloc(arena.base, ARENA_BLOCK_SIZE);
    ASAN_POISON(arena.base, ARENA_BLOCK_SIZE);
    if (!arena.base) {
        /* Hard fail on commit failure */
        sys_panic(STR("Failed to commit initial memory block: System may be out of memory"));
    }
    arena.committed = ARENA_BLOCK_SIZE;

    /* Arena should be 64k aligned */
    ASSERT(((u64)arena.base & 0xFFFF) == 0);

    return arena;
}

void arena_release(struct arena *arena)
{
    __prof;
    __proffree(arena->base);
    ASAN_UNPOISON(arena->base, arena->committed);
    sys_memory_release(arena->base);
}

/* NOTE: Application will exit if arena fails to commit memory */
void *_arena_push_bytes(struct arena *arena, u64 size, u64 align)
{
    ASSERT(align > 0);
    ASSERT(!arena->readonly);

    void *start = NULL;

    /* TODO: Remove this check? (Only here to avoid aligning when size = 0) */
    if (size > 0) {
        u64 aligned_start_pos = (arena->pos + (align - 1));
        aligned_start_pos -= aligned_start_pos % align;

        u64 new_pos = aligned_start_pos + size;
        if (new_pos > arena->committed) {
            __profscope(_arena_push_bytes_COMMIT);
            /* Commit new block(s) */
            u64 blocks_needed = (new_pos - arena->committed + ARENA_BLOCK_SIZE - 1) / ARENA_BLOCK_SIZE;
            u64 commit_bytes = blocks_needed * ARENA_BLOCK_SIZE;
            u64 new_capacity = arena->committed + commit_bytes;
            if (new_capacity > arena->reserved) {
                /* Hard fail if we overflow reserved memory for now */
                sys_panic(STR("Failed to commit new memory block: Overflow of reserved memory"));
            }
            void *commit_address = arena->base + arena->committed;
            if (!sys_memory_commit(commit_address, commit_bytes)) {
                /* Hard fail on memory allocation failure for now */
                sys_panic(STR("Failed to commit new memory block: System may be out of memory"));
            }
            __proffree(arena->base);
            __profalloc(arena->base, arena->committed + commit_bytes);
            ASAN_POISON(commit_address, commit_bytes);
            arena->committed += commit_bytes;
        }
        start = arena->base + aligned_start_pos;
        arena->pos = new_pos;
        ASAN_UNPOISON(start, (arena->base + arena->pos) - (u8 *)start);
    } else {
        start = arena->base + arena->pos;
    }

    return start;
}

/* Copies the memory from the source arena into the destination arena,
 * replacing old contents. Destination arena will be expanded if necessary. */
void arena_copy_replace(struct arena *dest, struct arena *src)
{
    arena_reset(dest);
    u64 data_size = src->pos;
    u8 *data_src = src->base;
    u8 *data_dest = _arena_push_bytes(dest, data_size, 1);
    MEMCPY(data_dest, data_src, data_size);
}

void arena_decommit_unused_blocks(struct arena *arena)
{
    ASSERT(!arena->readonly);
    u64 next_block_pos = ARENA_BLOCK_SIZE * ((arena->pos + (ARENA_BLOCK_SIZE - 1)) / ARENA_BLOCK_SIZE);
    if (arena->committed > next_block_pos) {
        u8 *decommit_start = arena->base + next_block_pos;
        u64 decommit_size = (arena->base + arena->committed) - decommit_start;
        sys_memory_decommit(decommit_start, decommit_size);
        arena->committed = next_block_pos;
    }
}

void arena_set_readonly(struct arena *arena)
{
    sys_memory_set_committed_readonly(arena->base, arena->committed);
#if RTC
    arena->readonly = true;
#endif
}