#pragma once

#include "base.h"

/*
===============================================================================
                               ARENA USAGE GUIDE
===============================================================================

OVERVIEW
--------
The arena allocator is a high-performance memory allocator designed for
predictable allocation patterns. It supports:

- Multiple chained memory blocks
- Pointer-style allocation with free-list reuse
- Stack-style allocation
- Optional Swap-back removal for unordered data
- Merging compatible arenas

Memory is committed in page-sized chunks. The arena never commits less than
the system page size and grows by allocating new blocks when needed.


CORE CONCEPTS
-------------

Arena Blocks
~~~~~~~~~~~~
An arena consists of one or more memory blocks linked together.

Each block contains:
- base_pos : The local starting position of a block (bytes)
- pos      : Global position (bytes)
- commit_size   : Total committed size (bytes)
- reserve_size  : Total usable size of the block (bytes)
- prev/next: Links to neighboring blocks
- the arena pointer points to the current block

The arena allocates from the current block where the global position is or from
the free list.

Blocks form a single logical address space:
    global_offset = pos


Global Position Model
~~~~~~~~~~~~~~~~~~~~~
All allocations are addressed using a global offset. This allows:
- Popping across block boundaries
- Cross-block swap-back operations
- Arena merging by linking blocks


ARENA CONFIGURATION
-------------------

An arena is configured usings the struct arena_params before creation.

Important parameters:
- push_size        : Fixed element size in bytes (0 = variable-size arena)
- allow_free_list  : Enables arena_free() and pointer mode, disabling it will
enable stack mode
- allow_swapback   : Enables arena_swapback_pop(), works in stack mode only
- max_nbre_blocks  : Maximum number of blocks (0 = unlimited)
- growth_policy    : How blocks grow (next created block will have the same size
or double the size)
- commit_policy    : How memory is committed (lazily if needed or commit all at
block creation)


ARENA CREATION AND DESTRUCTION
------------------------------

We create arenas using
    mem_arena *arena_create(arena_params *params);

Behavior:
- Create an arena and return the pointer to this arena or NULL if it fails

Requirements:
- The configuration of the arena needs to be injected as an argument using
arena_params

We destroy arenas using
    void *arena_destroy(mem_arena **arena_ptr);

Behavior:
- Return (void *)1 if it succeeds or NULL if it fails


ALLOCATION (arena_push)
-----------------------

We allocate using one function:
    void *arena_push(mem_arena **arena_ptr, u64 size, bool non_zero);

Behavior:
- zero can be set to true of false to zero the allocated memory or not
- If the current block is full and block_index < max_nbre_blocks, Grows into a
new block, setting the new block to arena->next then setting the arena pointer
to the new block
- The size can be set only if the arena has variable-size elements (push_size ==
0)
- Return the pointer of the pushed element if it succeeds or NULL if it fails

Variable-size allocation (allow_free_list == true, pointer mode only):

- Allocates 'size' bytes
- Only available in pointer mode with a free list
- Since the size of the elements is variable, when using the free list we loop
through all the elements until finding one with enough size, after allocation if
there is a remaining we store it as a new entry in the free list. The allocation
is slower than fixed-size arenas
- If there is not enough memory to push an element we try to create a new block
and push the element to it, if there is some memory remaining in the previous
block we add it to the free list
- max_nbre_blocks determines the behavior of the free list
    - If max_nbre_blocks > 0 we push directly to the arena until it's full then
we use the free list
    - If max_nbre_blocks = 0 (unlimited number of blocks) we use the free list
directly

Fixed-size allocation (allow_free_list can be true of false):

- Size is ignored, push_size defines element size (arena config)
- Available for both in stack and pointer modes
- Caller-provided size is ignored
- Required for swap-back correctness
- Faster and safer than variable-size mode
- If allow_free_list == true (pointer mode), uses the free list first then push

DEALLOCATION
------------
POINTER MODE (WITH A FREE LIST)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
- Popping requires the pointer that points to the element
- The only function allowed to pop is arena_free()

    void *arena_free(mem_arena **arena_ptr, u8 **ptr, u64 size);

Behavior:
- Stores (offset, size) in the free list
- Memory is reused on future allocations
- Sets the pointer to NULL
- The freed memory is not zeroed, we can zero the memory only when allocating
- If the element is last, behaves like a pop, decrementing pos, if the position
crosses to the previous arena block, set the arena pointer to arena->prev
- Return (void *)1 if it succeeds or NULL if it fails

Requirements:
- allow_free_list == true
- Correct element size for variable-size arenas

STACK MODE (HAS NO FREE LIST)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
- If allow_swapback = false, we can only pop with arena_pop_to()
- If allow_swapback = true, we can pop with arena_pop_to() and
arena_swapback_pop()

Pop to position:
    void *arena_pop_to(mem_arena **arena_ptr, u64 count);

- Removes x elements from the top of the stack
- Can remove data across blocks
- Pops the element by moving the position (pos)
- Does not zero memory we can zero the memory only when allocating
- If we pop to the previous arena block, set the arena pointer to arena->prev
- Return (void *)1 if it succeeds or NULL if it fails

Requirements:
- allow_free_list == false
- Fixed-size elements (push_size > 0)
- The number of elements to pop

A macro arena_pop(arena) is available to pop one element

void *arena_swapback_pop(mem_arena **arena_ptr, u64 index);

- The top of the stack element is copied into the removed slot, the top of the
stack is then reduced by one element
- Works across multiple blocks
- Return (void *)1 if it succeeds or NULL if it fails
- Swapping back the last element will only pop it

Cases:
- Middle of block        -> swapback
- End of non-last block  -> swapback
- Last element overall   -> pop only

Requirements:
- allow_free_list == false
- Fixed-size elements (push_size > 0)
- allow_swapback == true


CLEARING THE ARENA
------------------

    void *arena_clear(mem_arena **arena_ptr);

- Resets pos for all blocks
- Keeps all committed memory
- Does NOT zero memory
- Resets the arena pointer to the first arena block
- Resets the free-list
- Return (void *)1 if it succeeds or NULL if it fails


MERGING ARENAS
--------------

    mem_arena *arena_merge(mem_arena **dst_ptr, mem_arena **src_ptr);

Behavior:
Merges arena B into arena A by:
- Updating the pos and base_pos and block_index of all the blocks of B
- Linking the two arenas by setting arena->prev of the first block of B to the
last block of A
- If there is empty blocks in arena A, append them to the end of arena B
- Resets the arena pointer to arena B
- Merges the free list arenas if available
- max_nbre_blocks is summed
- Return the pointer to the new arena of NULL if it fails

Requirements:
- Configurations must match exactly

HELPER FUNCTIONS
------------------

mem_arena *arena_block_from_pos(mem_arena *last, u64 global_pos);
mem_arena *arena_block_from_index(mem_arena *last, u64 index);
mem_arena *arena_block_from_ptr(mem_arena *last, u8 *ptr);
u64 arena_pos_from_ptr(mem_arena *arena, void *ptr);
void *arena_ptr_from_pos(mem_arena *arena, u64 global_pos);
void *arena_ptr_from_index(mem_arena *arena, u64 index);
===============================================================================
*/

#define ARENA_HEADER_SIZE (sizeof(mem_arena))
#define ARENA_ALIGN (sizeof(void *))

// arena config
typedef enum arena_growth_policy {
  ARENA_GROWTH_NORMAL = 0, // grow by fixed block size
  ARENA_GROWTH_DOUBLE      // double block size on each growth
} arena_growth_policy;

typedef enum arena_commit_policy {
  ARENA_COMMIT_LAZY = 0, // commit pages on demand
  ARENA_COMMIT_FULL      // commit entire reserve at creation
} arena_commit_policy;

typedef struct arena_params {
  u64 reserve_size; // size of one arena block
  u64 commit_size;  // initial commit size
  u64 align; // allocation alignment, 0 to disable and ARENA_ALIGN to align
             // according to architecture

  // Element size rules:
  // - stack mode        : push_size > 0 (mandatory)
  // - pointer fixed     : push_size > 0
  // - pointer variable  : push_size == 0
  u64 push_size;

  struct mem_arena *free_list;
  b32 allow_free_list; // pointer mode if true

  b32 allow_swapback; // stack mode only

  arena_growth_policy growth_policy;
  arena_commit_policy commit_policy;

  u32 max_nbre_blocks; // 0 = unlimited
} arena_params;

typedef struct arena_free_node {
  u64 offset; // offset from arena base
  u64 size;   // size of freed block
} arena_free_node;

// arena definition struct
typedef struct mem_arena {
  // block chaining
  struct mem_arena *prev;
  struct mem_arena *next; // valid only on root arena

  // positions
  u64 base_pos; // of selected block
  u64 pos;      // global pos

  // memory limits
  u64 reserve_size;
  u64 commit_size;
  u64 commit_pos;

  // configuration
  u64 align;
  u64 push_size;

  // Pointer mode only
  struct mem_arena *free_list;
  b32 allow_free_list;

  // Stack mode only
  b32 allow_swapback;

  arena_growth_policy growth_policy;
  arena_commit_policy commit_policy;
  u32 max_nbre_blocks;
  u32 block_index;

} mem_arena;

typedef struct mem_arena_temp {
  mem_arena *arena;
  u64 pos;
} mem_arena_temp;

// helper functions
mem_arena *arena_block_from_pos(mem_arena *last, u64 global_pos);
mem_arena *arena_block_from_index(mem_arena *last, u64 index);
mem_arena *arena_block_from_ptr(mem_arena *last, u8 *ptr);
u64 arena_pos_from_ptr(mem_arena *arena, void *ptr);
void *arena_ptr_from_pos(mem_arena *arena, u64 global_pos);
void *arena_ptr_from_index(mem_arena *arena, u64 index);

// arena core functions
// creation / destruction
mem_arena *arena_create(arena_params *params);
void *arena_destroy(mem_arena **arena_ptr);

// allocation
void *arena_push(mem_arena **arena_ptr, u64 size, bool non_zero);

// pointer mode only
// - fixed-size arena   : size is ignored
// - variable-size arena: size is mandatory
void *arena_free(mem_arena **arena_ptr, u8 **ptr, u64 size);

// stack mode only
void *arena_pop_to(mem_arena **arena_ptr, u64 count);
void *arena_swapback_pop(mem_arena **arena_ptr, u64 index);

// utilities
void *arena_clear(mem_arena **arena_ptr);

mem_arena *arena_merge(mem_arena **dst_ptr, mem_arena **src_ptr);

// temp arenas
/* create a temporary arena on top of an existing arena, it:
 * takes an existing arena, marks it's position, allocate beyond this position
 * and free it until the marked position when not needed*/
mem_arena_temp arena_temp_begin(mem_arena *arena);
void arena_temp_end(mem_arena_temp temp);

mem_arena_temp arena_scratch_get(mem_arena **conflicts, u32 num_conflicts);
void arena_scratch_release(mem_arena_temp scratch);

#if defined(_MSC_VER)
#define THREAD_LOCAL __declspec(thread)
#else
#define THREAD_LOCAL __thread
#endif

// Helpers
// Pointer mode only
/* Fixed-size arena: size is implicit, can pass 0 */
#define PUSH_STRUCT(arena, T)                                                  \
  ((T *)arena_push((arena), 0, true)) // Zeroes the memory
#define PUSH_STRUCT_NZ(arena, T) ((T *)arena_push((arena), 0, false))

#define PUSH_ARRAY(arena, T, n) ((T *)arena_push((arena), 0, true))
#define PUSH_ARRAY_NZ(arena, T, n) ((T *)arena_push((arena), 0, false))

/* Variable-size arena helpers: REQUIRE explicit size via arena_push() */
#define ARENA_PUSH(arena, size) arena_push((arena), (size), true)
#define ARENA_PUSH_NZ(arena, size) arena_push((arena), (size), false)

#define arena_pop(arena_ptr) arena_pop_to((arena_ptr), 1)