#include <stdio.h>
#include "malunal/allocators/arena.h"

#define assert_equals_result(condition) \
  if (!(condition)) {                    \
    fprintf(                            \
      stderr,                           \
      format,                           \
      __FILE__,                         \
      __LINE__,                         \
      MALUNAL_TOSTRING(condition)       \
    );                                  \
    return (allocation_result_t) {      \
      .threw = 1,                       \
      .error = ALLOCATION_ERROR_FAILURE \
    };                                  \
  }

#define assert_equals(condition)  \
  if (!(condition)) {             \
    fprintf(                      \
      stderr,                     \
      format,                     \
      __FILE__,                   \
      __LINE__,                   \
      MALUNAL_TOSTRING(condition) \
    );                            \
    return 0;                     \
  }


const malunal_cstr_t format = "[arena_allocator(%s:%d)]: %s failed\n";

const malunal_size_t page_bytes     = 4096;
const malunal_size_t pre_alloc      = 16;
const malunal_size_t init_bytes     = sizeof(malunal_int32_t) * 4;
const malunal_size_t final_bytes    = sizeof(malunal_int32_t) * 8;
const malunal_size_t init_with_pre  = pre_alloc + init_bytes;
const malunal_size_t all_bytes      = pre_alloc + init_bytes + final_bytes;

static
allocation_error_t
test_arena_initialize(
  allocator_mptr_t arena,
  allocator_mptr_t upstream
) {
  assert_equals(allocator_initialize(arena, upstream, page_bytes) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  assert_equals(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  // Arena checks.
  assert_equals(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == pre_alloc);

  assert_equals(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);
  return ALLOCATION_ERROR_SUCCESS;
}

static
allocation_result_t
test_arena_acquires(allocator_mptr_t arena) {
  allocation_result_t result = allocator_acquire(arena, init_bytes);
  if (result.threw)
    return result;

  // Upstream checks.
  allocator_mptr_t upstream;
  assert_equals_result(allocator_upstream(arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals_result(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == page_bytes);

  assert_equals_result(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == page_bytes);

  assert_equals_result(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == 1);

  assert_equals_result(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == 0);

  // Arena checks.
  assert_equals_result(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == page_bytes);

  assert_equals_result(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == init_with_pre);

  assert_equals_result(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == 1);

  assert_equals_result(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals_result(count == 0);
  return result;
}

static int
test_arena_reacquire(
  allocator_mptr_t arena,
  malunal_mptr_t address
) {
  malunal_int32_t  index;
  malunal_int32_t* vector;

  vector = address;
  for (index = 0; index < 4; index++)
    vector[index] = index + 1;

  allocation_result_t result = allocator_reacquire(
    arena,
    address,
    init_bytes,
    final_bytes
  );

  if (result.threw)
    return result.error;

  allocator_mptr_t upstream;
  assert_equals(allocator_upstream(arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  assert_equals(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  // Arena checks.
  assert_equals(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == all_bytes);

  assert_equals(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);

  assert_equals(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);
  return ALLOCATION_ERROR_SUCCESS;
}

static int
test_arena_release(
  allocator_mptr_t arena,
  malunal_mptr_t   address
) {
  malunal_int32_t  index;
  malunal_int32_t* vector;

  vector = address;
  for (index = 0; index < 4; index++)
    if (vector[index] != index + 1)
      return 1;

  allocation_error_t relexc =
    allocator_release(arena, address, final_bytes);

  if (relexc != ALLOCATION_ERROR_SUCCESS)
    return relexc;

  allocator_mptr_t upstream;
  assert_equals(allocator_upstream(arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  assert_equals(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  // Arena checks.
  assert_equals(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == all_bytes);

  assert_equals(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);

  assert_equals(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);
  return ALLOCATION_ERROR_SUCCESS;
}

static int
test_arena_reset(allocator_mptr_t arena) {
  arena_reset(arena);

  allocator_mptr_t upstream;
  assert_equals(allocator_upstream(arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  assert_equals(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  // Arena checks.
  assert_equals(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == page_bytes);

  assert_equals(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);

  assert_equals(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);
  return ALLOCATION_ERROR_SUCCESS;
}

static int
test_arena_finalize(allocator_mptr_t arena) {
  allocator_finalize(arena);

  allocator_mptr_t upstream;
  assert_equals(allocator_upstream(arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  malunal_size_t count;
  assert_equals(allocator_reserved_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_allocated_bytes(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_acquires_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  assert_equals(allocator_releases_count(upstream, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 1);

  // Arena checks.
  assert_equals(allocator_reserved_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_allocated_bytes(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 0);

  assert_equals(allocator_acquires_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);

  assert_equals(allocator_releases_count(arena, &count) == ALLOCATION_ERROR_SUCCESS);
  assert_equals(count == 2);
  return ALLOCATION_ERROR_SUCCESS;
}

int test_arena_allocator() {
  allocator_t upstream = platform_allocator();
  allocator_t arena    = arena_allocator();
  assert_equals(test_arena_initialize(&arena, &upstream) == ALLOCATION_ERROR_SUCCESS);

  allocation_result_t result =
    test_arena_acquires(&arena);

  return !(
    !result.threw &&
    test_arena_reacquire(&arena, result.address) &&
    test_arena_release(&arena, result.address)   &&
    test_arena_reset(&arena)                     &&
    test_arena_finalize(&arena)
  );
}