#ifndef FOR_TESTING
#define FOR_TESTING
#endif

#include <stdlib.h>

#include "arch.h"
#include "test_harness.c"
#include "common.h"
#include "mem.h"

extern uint32_t* halloc_start;
static void wipeout_halloc()
{
  memset(halloc_start, 0, 1024);
  halloc_start = NULL;
}

struct TEST_STRUCT {
  uint32_t array[3];
};

struct TEST_STRUCT2 {
  uint32_t array[10];
};


static struct TEST_STRUCT* new_test_struct()
{
  struct TEST_STRUCT* ret = halloc(sizeof(struct TEST_STRUCT));
  
  ret->array[0] = 1;
  ret->array[1] = 2;
  ret->array[2] = 3;

  return ret;
}

static struct TEST_STRUCT2* new_test_struct2()
{
  struct TEST_STRUCT2* ret = halloc(sizeof(struct TEST_STRUCT2));
  
  for (int i = 0; i < 10; ++ i) {
    ret->array[i] = i;
  }

  return ret;
}

#define ASSERT_CHAIN(t1, t2) \
  ASSERT_EQ(V(t1) + sizeof(*t1) + 4, V(t2))

TEST(memory, halloc)
{

#define V(x) ((void*)(x))
  struct TEST_STRUCT* test1 = new_test_struct();
  struct TEST_STRUCT2* test2 = new_test_struct2();
  struct TEST_STRUCT* test3 = new_test_struct();
  struct TEST_STRUCT2* test4 = new_test_struct2();
  struct TEST_STRUCT2* test5 = new_test_struct2();

  ASSERT_TRUE(V(test1) != V(test2));
  ASSERT_TRUE(V(test2) != V(test3));
  ASSERT_TRUE(V(test3) != V(test1));
  ASSERT_TRUE(V(test2) != V(test5));
  ASSERT_TRUE(V(test4) != V(test5));

  ASSERT_CHAIN(test1, test2);
  ASSERT_CHAIN(test2, test3);
  ASSERT_CHAIN(test3, test4);
  ASSERT_CHAIN(test4, test5);

  wipeout_halloc();

  return 0;
}

struct UNEVEN_STRUCT {
  uint8_t arr[5];
};

struct UNEVEN_STRUCT* new_uneven_struct()
{
  struct UNEVEN_STRUCT* ret = halloc(sizeof(struct UNEVEN_STRUCT));
  
  ret->arr[0] = 1;
  ret->arr[1] = 2;
  ret->arr[2] = 3;
  ret->arr[3] = 4;
  ret->arr[4] = 5;

  return ret;
}

#define size_for(n) \
  (((n) / 4) + ((n) % 4 != 0))

TEST(memory, uneven_halloc)
{
  struct UNEVEN_STRUCT* test1 = new_uneven_struct();
  struct UNEVEN_STRUCT* test2 = new_uneven_struct();

  ASSERT_EQ(V(test1) + 12, test2);

  wipeout_halloc();

  return 0;
}

TEST(memory, halloc_free)
{
  struct TEST_STRUCT* test1 = new_test_struct();
  struct TEST_STRUCT2* test2 = new_test_struct2();
  struct TEST_STRUCT* test3 = new_test_struct();
  struct TEST_STRUCT2* test4 = new_test_struct2();
  struct TEST_STRUCT2* test5 = new_test_struct2();

  hfree(test2);
  hfree(test4);
  hfree(test3);
  hfree(test1);
  hfree(test5);

  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);

  test1 = new_test_struct();
  test2 = new_test_struct2();
  test3 = new_test_struct();
  test4 = new_test_struct2();
  test5 = new_test_struct2();

  hfree(test1);
  hfree(test3);
  hfree(test2);
  hfree(test4);
  hfree(test5);

  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);

  test1 = new_test_struct();
  test2 = new_test_struct2();
  test3 = new_test_struct();
  test4 = new_test_struct2();
  test5 = new_test_struct2();

  hfree(test4);
  hfree(test3);
  hfree(test1);
  hfree(test2);
  hfree(test5);

  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);

  wipeout_halloc();

  return 0;
}

TEST(memory, halloc_free_alloc2)
{
  struct TEST_STRUCT2* test1 = new_test_struct2();
  struct TEST_STRUCT2* test2 = new_test_struct2();

  hfree(test1);

  struct TEST_STRUCT* test3 = new_test_struct();
  struct TEST_STRUCT* test4 = new_test_struct();

  ASSERT_EQ(debug_halloc_get_next_ptr(test3), V(test4));

  ASSERT_EQ(
      // There should be a free block after test4.
      debug_halloc_get_next_ptr(debug_halloc_get_next_ptr(test4)),
      V(test2));

  ASSERT_EQ(
      // There should be a free block after test4.
      debug_halloc_get_prev_ptr(debug_halloc_get_prev_ptr(test2)),
      V(test4));

  char buf[1024];
  if (debug_halloc_assert_consistency(buf, 1024)) {
    fprintf(stderr, "Consistency check failed.\n");
    fprintf(stderr, buf);
    ASSERT_TRUE(false);
  }

  return 0;
}

TEST(memory, consistency_stress)
{
  if (halloc_start) {
    wipeout_halloc();
  }

  int i;
  void* allocd[500] = { 0 };
  char buf[1024];

  for (i = 0; i < 500; ++ i) {
    size_t nalloc = rand() % 20;
    allocd[i] = halloc(nalloc);

    ASSERT_TRUE(allocd[i]);
    size_t idx = rand() % 500;

    hfree(allocd[idx]);
    allocd[idx] = NULL;

    idx = rand() % 500;
    hfree(allocd[idx]);
    allocd[idx] = NULL;

    if (debug_halloc_assert_consistency(buf, 1024)) {
      fprintf(stderr, "Consistency check failed. (At index %d) nalloc=%lu\n", i, nalloc);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }
  }

  for(i = 0; i < 500; ++ i) {
    hfree(allocd[i]);
  }

  if (debug_halloc_assert_consistency(buf, 1024)) {
    fprintf(stderr, "Consistency check failed. (Final)\n");
    fprintf(stderr, buf);
    ASSERT_TRUE(false);
  }
  ASSERT_EQ((*halloc_start >> 1) * 4, MAX_HEAP_SIZE - 4);

  return 0;
}

TEST(memory, halloc_free_alloc)
{
  new_test_struct();
  struct TEST_STRUCT2* test2 = new_test_struct2();
  new_test_struct();
  struct TEST_STRUCT2* test4 = new_test_struct2();
  new_test_struct2();

  hfree(test4);

  struct TEST_STRUCT2* test6 = new_test_struct2();

  // test_6 should have been allocated in test_4's spot.
  ASSERT_EQ(test6, test4); 

  hfree(test2);
  struct TEST_STRUCT* test7 = new_test_struct();
  struct TEST_STRUCT* test8 = new_test_struct();

  // Test 2 was large enough to accomodate 3 smaller structs.
  ASSERT_EQ(V(test7), V(test2));
  ASSERT_EQ(V(test8), V(test2) + sizeof(*test7) + 4);

  return 0;
}