#include "Arena.h"
#include "Arena_Windows.cpp"

// For arrays, use `Array`, which is backed by the general purpose allocator
// or use `ArenaArray` if you need to expand the size to an unknown size.
void* arena_allocator_proc (Allocator_Mode mode, s64 requested_size, s64 old_size, void* old_memory, void* allocator_data) {
  Arena* arena = (Arena*)allocator_data;
  Assert(arena != nullptr);
  switch (mode) {
    case Allocator_Mode::ALLOCATE: {
      return arena_alloc(arena, requested_size);
    } break;
    case Allocator_Mode::RESIZE: {
      Assert(false); // DO NOT USE RESIZE WITH ARENAS! :ArenaResizing
      // Or maybeeee...
      // Resize should check if current_point matches the end of the old allocation?
      // and resize accordingly + pass back same pointer.
      void* result = arena_alloc(arena, requested_size);
      s64 size_to_copy = old_size < requested_size ? old_size : requested_size;
      if (result && size_to_copy) {
        memcpy(result, old_memory, size_to_copy);
      }
      return result;
    } break;
    case Allocator_Mode::DEALLOCATE:
      return nullptr; // unused
    break;
  }
  
  return nullptr;
}

// See arena_delete for decommissioning of arenas initialized with arena_init
void arena_init (Arena* arena, Arena_Reserve new_reserve, s32 default_commit_page_count) {
  Assert(arena != nullptr);
  
  s64 commit_size = default_commit_page_count * PLATFORM_MEMORY_PAGE_SIZE;
  Assert(commit_size >= 0);
  
  if (!is_valid(arena)) {
    arena->reserve_size = new_reserve;
    
    platform_init(arena, reserve_size(arena));
    arena->current_point = arena_start(arena);
    arena->first_uncommitted_page = arena->memory_base;
    if (arena_commit_first_pages(arena, commit_size)) {
      arena->initial_commit_page_count = default_commit_page_count;
      // arena_lock_pages(arena, arena->memory_base, commit_size);
    }
    
    arena->alignment = ARENA_DEFAULT_ALIGNMENT;
    arena->flags = (arena->flags | (~Arena_Flags::Is_Bootstrapped)); // unset is bootstrapped, if present.
  }
}

bool arena_commit_first_pages (Arena* arena, s64 commit_size, s64 start_offset) {
  Assert(commit_size > 0); // no point calling this without commit size
  Assert(start_offset >= 0); // should NOT be negative.
  
  if (commit_size <= 0 || start_offset < 0) {
    return false;
  }

  u8* final_address = Align<u8*>(arena->memory_base + commit_size + start_offset, 
                              arena->alignment);
  s64 commit_size_final = (s64)(final_address - arena->memory_base);
  
  if (commit_size_final > reserve_size(arena)) {
    Assert(false); // Commit size exceeds reserve size
    return false;
  }
  
  if (final_address > arena->first_uncommitted_page) {
    extend_committed_pages(arena, final_address);
  }
  
  return true;
}

Arena* bootstrap_arena (Arena_Reserve new_reserve, s32 default_commit_page_count) {
  s64 commit_size = default_commit_page_count * PLATFORM_MEMORY_PAGE_SIZE;
  Assert(commit_size <= reserve_size(new_reserve));
  
  Arena new_arena;
  arena_init(&new_arena, new_reserve, default_commit_page_count);
  Arena* arena_ptr = (Arena*)new_arena.memory_base;
  if (arena_ptr == nullptr) { Assert(false); // Assert in debug mode!
    return nullptr;
  }
  
  memcpy(arena_ptr, &new_arena, sizeof(Arena));  
  arena_ptr->current_point = arena_start(arena_ptr);
  
  arena_set_bootstrap_flag(arena_ptr);
  
  return arena_ptr;
}

void arena_clear_flags (Arena* arena) { arena->flags = Arena_Flags::None; }
void arena_set_bootstrap_flag (Arena* arena) { arena->flags |= Arena_Flags::Is_Bootstrapped; }
bool arena_is_bootstrapped (Arena* arena) { return (bool)(arena->flags & Arena_Flags::Is_Bootstrapped); }
// void arena_set_secure_flag (Arena* arena) { arena->flags |= Arena_Flags::Secure_Arena; }

void arena_reset_keeping_memory (Arena* arena) {
  if (!is_valid(arena)) return;
  
  arena->current_point = arena_start(arena);
}

void arena_reset (Arena* arena) {
  if (!is_valid(arena)) return;
  
  arena->current_point = arena_start(arena);
  
  free_pages_down_to(arena, arena->initial_commit_page_count);
}

void arena_reset_overwriting_memory (Arena* arena, Memory_Wipe_Function wipe_function) {
  if (!is_valid(arena)) return;
  Assert(wipe_function != nullptr);
  if (wipe_function == nullptr) return;
  wipe_function(arena_start(arena), (u64)(arena->current_point - arena_start(arena)));
  
  arena_reset(arena);
}

void* arena_alloc (Arena* arena, s64 byte_count) {
  Assert(arena != nullptr);
  
  if (!arena->memory_base) { arena_init(arena, arena->reserve_size); }
  Assert(is_valid(arena));
  
  u8* result = Align<u8*>(arena->current_point, arena->alignment);
  u8* result_end = result + byte_count;
  
  if (result_end > arena->first_uncommitted_page) {
    if (result_end > arena_address_limit(arena)) {
      printf("[Error] Failed to allocate because Arena is full and cannot expand!\n");
      Assert(false);  // Failed to allocate because arena is full and cannot expand!
    } else {
      extend_committed_pages(arena, result_end);
    }
  }
  
  arena->current_point = result_end;
  
  return result;
}

u8* arena_start (Arena* arena) {
  // This assumes every arena is bootstrapped, so there may be some 
  // wastage/cache misses caused by this!
  return Align(arena->memory_base + sizeof(Arena), ARENA_DEFAULT_ALIGNMENT);
}

u8* arena_address_limit (Arena* arena) {
  return (arena->memory_base + reserve_size(arena));
}

bool is_valid (Arena* arena) {
  return (arena != nullptr) && (arena->memory_base != nullptr);
}

s64 reserve_size (Arena* arena) {
  return reserve_size(arena->reserve_size);
}

s64 reserve_size (Arena_Reserve ar) {
  switch (ar) {
    case Arena_Reserve::Size_64K: return KB(64); 
    case Arena_Reserve::Size_2M:  return MB(2);
    case Arena_Reserve::Size_64M: return MB(64);
    case Arena_Reserve::Size_2G:  return GB(2);
    case Arena_Reserve::Size_64G: return GB(64);
    case Arena_Reserve::Size_2T:  return TB(2);
  }
  
  Assert(false); // This should not happen.
  
  return 0;
}

Arena_Reserve next_reserve_size (s64 size) {
  for (u8 i = 0; i < Arena_Reserve_Count; i += 1) {
    if (size <= Arena_Sizes[i]) {
      return (Arena_Reserve)i;
    }
  }
  
  return Arena_Reserve::Size_64T;
}

// arena_usage_bytes is kinda pointless tbh.
s64 arena_usage_bytes (Arena* arena) { return (s64)(arena->current_point - arena_start(arena)); }
s64 arena_usage_committed_bytes (Arena* arena) { return (s64)(arena->first_uncommitted_page - arena->memory_base); }

// for arena details, I need to setup my string builder first.
Allocator allocator (Arena* arena) {
  return { arena_allocator_proc, arena };
}

struct Auto_Reset {
  Arena* arena;
  u8* starting_point;

  Auto_Reset(Arena* arena) {
    Assert(is_valid(arena));
    this->arena = arena;
    this->starting_point = arena->current_point;
  }
  
  // #TODO: Implement with ExpandableArena
  // Auto_Reset(ExpandableArena* arena_ex) {
  //   Auto_Reset((Arena*)arena_ex);
  // }
  
  ~Auto_Reset() {
    arena->current_point = starting_point;
  }
};

struct Auto_Release {
  bool is_expandable;
  ExpandableArena* arena_ex;
  Arena* arena;
  u8* starting_point;

  Auto_Release(ExpandableArena* arena_ex) {
    this->arena_ex = arena_ex;
    this->arena    = arena_ex->current;
    this->starting_point = arena_ex->current->current_point;
    this->is_expandable = true;
  }

  Auto_Release(Arena* arena) {
    this->arena_ex = nullptr;
    this->arena = arena;
    this->starting_point = arena->current_point;
    this->is_expandable = false;
    Assert(is_valid(arena));
  }
  
  ~Auto_Release() {
    if (is_expandable) {
      arena_reset_to(arena_ex, arena, starting_point);
    } else {
      arena->current_point = starting_point;
      free_pages_down_to(arena, arena->initial_commit_page_count);
    }
  }
};