// NOTE: Includes ==================================================================================
#define NOMINMAX
#define WIN32_MEAN_AND_LEAN
#include <Windows.h>

#include <assert.h>
#include <stdint.h>
#include <stdio.h>

// NOTE: Types =====================================================================================
typedef uint16_t U16;
typedef uint32_t U32;
typedef uint64_t U64;

#define MIN(a, b) (a) < (b) ? (a) : (b)
#define MAX(a, b) (a) > (b) ? (a) : (b)

struct OS
{
    U32 page_size;
    U32 alloc_granularity;
};
OS g_os;

struct ArenaBlock
{
    void       *data;
    U64         pos, commit, cap;
    ArenaBlock *next, *prev;
};

struct Arena
{
    ArenaBlock *first, *last;
    U32         flags;
};

enum ArenaFlag
{
    ArenaFlag_Chain              = 1 << 0,
    ArenaFlag_ReserveCommit      = 1 << 1,
    ArenaFlag_ChainReserveCommit = ArenaFlag_Chain | ArenaFlag_ReserveCommit,
};

#define KILOBYTES(val) 1024ULL * val
#define MEGABYTES(val) 1024ULL * KILOBYTES(val)
#define GIGABYTES(val) 1024ULL * MEGABYTES(val)
#define TERABYTES(val) 1024ULL * GIGABYTES(val)

// NOTE: API =======================================================================================
Arena *ArenaAlloc(U64 cap, uint32_t flags);
void   ArenaRelease(Arena *arena);
U64    ArenaPos(Arena *arena);

#define ArenaPushN(arena, T, size) (T *)ArenaPush(arena, sizeof(T) * size, alignof(T))
void  *ArenaPushNoZero(Arena *arena, U64 size, U64 alignment);
void  *ArenaPush(Arena *arena, U64 size, U64 alignment);
void   ArenaPushAlignment(Arena *arena, U64 alignment);

void   ArenaPopTo(Arena *arena, U64 pos);
void   ArenaPop(Arena *arena, U64 size);
void   ArenaClear(Arena *arena);

#define IsPowerOfTwo(value, pot) ((((uintptr_t)value) & ((uintptr_t)pot - 1)) == 0)
#define AlignUpPowerOfTwo(value, pot) (((uintptr_t)(value) + ((uintptr_t)(pot) - 1)) & ~((uintptr_t)(pot) - 1))
#define AlignDownPowerOfTwo(value, pot) (((uintptr_t)value) & (~((uintptr_t)pot - 1)))
#define ArenaBlockMetadataSize(block) sizeof(ArenaBlock) + (block->next == block->prev ? sizeof(Arena) : 0)

template <typename T>
struct Array
{
    T      *data;
    size_t  size;
    size_t  cap;
    Arena  *arena;

    static Array Init(U64 capacity);

    T       *begin()       { return data; }
    T       *end()         { return data + size; }
    T const *begin() const { return data; }
    T const *end()   const { return data + size; }

    T *PushBack(const T& item);
};

template <typename T>
Array<T> Array<T>::Init(U64 capacity)
{
    Array<T> result = {};
    result.arena    = ArenaAlloc(capacity * sizeof(T), ArenaFlag_ReserveCommit);
    result.cap      = capacity;
    result.data     = (T *)ArenaPush(result.arena, 0 /*size*/, alignof(T));
    return result;
}

template <typename T>
T *Array<T>::PushBack(const T& item)
{
    T* result = ArenaPushN(arena, T, 1);
    if (result) {
        size++;
        *result = item;
    }
    return result;
}

// NOTE: Implementation ============================================================================
Arena *ArenaAlloc(U64 initial_cap, U32 flags)
{
    assert((flags & ~ArenaFlag_ChainReserveCommit) == 0);

    // NOTE: Allocate memory
    U64 metadata_size = sizeof(Arena) + sizeof(ArenaBlock);
    U64 alloc_size    = AlignUpPowerOfTwo(metadata_size + initial_cap, g_os.page_size);
    U32 os_flags      = MEM_RESERVE | ((flags & ArenaFlag_ReserveCommit) ? 0 : MEM_COMMIT);
    char *ptr         = (char *)VirtualAlloc(nullptr, alloc_size, os_flags, PAGE_READWRITE);
    char const *end   = ptr + alloc_size;

    U64 real_commit_size = 0;
    if ((os_flags & MEM_COMMIT)) {
        real_commit_size = alloc_size;
    } else {
        real_commit_size = AlignUpPowerOfTwo(metadata_size, g_os.page_size);
        VirtualAlloc(ptr, real_commit_size, MEM_COMMIT, PAGE_READWRITE);
    }

    // NOTE: Create arena and block
    Arena *result      = (Arena *)ptr;
    result->flags      = flags;
    ptr               += sizeof(*result);
    ArenaBlock *block  = (ArenaBlock *)ptr;
    ptr               += sizeof(*block);
    block->data        = ptr;
    block->cap         = end - ptr;
    block->commit      = real_commit_size - metadata_size; // NOTE: Raw commit size includes block metadata
    assert(IsPowerOfTwo(block, alignof(ArenaBlock)));      // NOTE: Enforce alignment of block

    // NOTE: Setup linked list of blocks
    block->next   = block->prev = block;
    result->first = result->last = block->next;
    return result;
}

void ArenaRelease(Arena *arena)
{
    if (!arena)
        return;
    for (ArenaBlock *block = arena->first, *next = block->next; block->next != block->prev; block = next)
        VirtualFree(block, 0, MEM_RELEASE);
    VirtualFree(arena, 0, MEM_RELEASE);
}

U64 ArenaPos(Arena *arena)
{
    U64 result = (arena && arena->last) ? arena->last->pos : 0;
    return result;
}

void *ArenaPushNoZero(Arena *arena, U64 size, U64 alignment)
{
    void *result = nullptr;
    if (!arena)
        return result;

    assert(IsPowerOfTwo(alignment, alignment)); // Enforce PoT
    U64 new_block_pos = 0;
    for (;;) {
        new_block_pos = AlignUpPowerOfTwo(arena->last->pos, alignment) + size;
        if (new_block_pos <= arena->last->cap)
            break;

        if ((arena->flags & ArenaFlag_Chain) == 0)
            return result;

        U64 alloc_size    = AlignUpPowerOfTwo(sizeof(ArenaBlock) + size, g_os.page_size);
        ArenaBlock *block = (ArenaBlock *)VirtualAlloc(nullptr, alloc_size, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
        if (!block)
            return result;

        block->cap        = alloc_size - sizeof(*block);
        block->data       = (char *)block + sizeof(*block);
        block->prev       = arena->last;
        block->next       = arena->last->next;
        block->prev->next = block;
        block->next->prev = block;
        arena->last       = block;
        block->commit     = block->cap;
        assert(((uintptr_t)block & alignof(ArenaBlock) - 1) == 0); // Enforce alignment of block
    }

    // NOTE: Align PoT & divvy out the pointer
    ArenaBlock *block = arena->last;
    block->pos        = new_block_pos;
    result            = (char *)block->data + (block->pos - size);

    // NOTE: Commit pages
    if (block->commit <= block->pos) {
        U64 commit_size  = AlignUpPowerOfTwo(block->pos - block->commit, g_os.page_size);
        void *commit_ptr = (void *)AlignUpPowerOfTwo((char *)block->data + block->commit, g_os.page_size);
        VirtualAlloc(commit_ptr, commit_size, MEM_COMMIT, PAGE_READWRITE);
        block->commit   += commit_size;
    }

    assert(IsPowerOfTwo(result, alignment));
    return result;
}

void *ArenaPush(Arena *arena, U64 size, U64 alignment)
{
    void *result = ArenaPushNoZero(arena, size, alignment);
    if (result)
        memset(result, 0, size);
    return result;
}

void ArenaPopTo(Arena *arena, U64 pos)
{
    if (arena) {
        assert(pos <= arena->last->pos);
        arena->last->pos = pos;
    }
}

void ArenaPop(Arena *arena, U64 size)
{
    if (arena) {
        assert(size <= arena->last->pos);
        arena->last->pos -= size;
    }
}

void ArenaClear(Arena *arena)
{
    if (arena)
        arena->last->pos = 0;
}

// NOTE: Main ============================================================================
int main(int, char**)
{
    SYSTEM_INFO info = {};
    GetSystemInfo(&info);
    g_os.page_size         = info.dwPageSize;
    g_os.alloc_granularity = info.dwAllocationGranularity;
    printf("Page Size: %u bytes\n", g_os.page_size);
    printf("Allocation Granularity: %u bytes\n", g_os.alloc_granularity);

    Arena *arena = ArenaAlloc(g_os.page_size - 8, ArenaFlag_ChainReserveCommit);
    {
        U64 u64s_in_a_page = g_os.page_size / sizeof(U64);
        U64 *u64s_array  = ArenaPushN(arena, U64, u64s_in_a_page);
        assert(u64s_array);
        memset(u64s_array, 0xFA, u64s_in_a_page * sizeof(U64));
    }

    {
        U64 u64s_in_straddled_page = (g_os.page_size / sizeof(U64)) + 1;
        U64 *u64s_array  = ArenaPushN(arena, U64, u64s_in_straddled_page);
        assert(u64s_array);
        memset(u64s_array, 0xFB, u64s_in_straddled_page * sizeof(U64));
    }

    {
        Array<U64> array = Array<U64>::Init(TERABYTES(1) / sizeof(U64));
        for (size_t index = 0; index < 100'000; index++)
            array.PushBack(1ULL);
        assert(array.size == 100'000);
        assert(array.cap == TERABYTES(1) / sizeof(U64));
        memset(array.data, 0xFC, array.size * sizeof(*array.data));
    }

    ArenaPopTo(arena, 0);
    return 0;
}