#ifndef FOR_TESTING
#define FOR_TESTING
#endif

#include <stdlib.h>

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

struct TEST_STRUCT {
  uint32_t array[3];
};

struct TEST_STRUCT2 {
  uint32_t array[10];
};

/* Copy of the node structure. */
typedef struct KALLOC_NODE {
  union {
    uint32_t header;
    struct {
      /* Is this memory block currently in use (hasn't been kfree'd) */
      bool used : 1;
      /* Number of words allocated. Does not include the header. */
      uint16_t size : 12;
      /* The location of the previous block (in WORDS from offest) */
      uint16_t prev : 12;
      uint8_t canary : 7;
    } PACKED;
  };

  uint8_t mem[]; /* The memory to use. */
} kalloc_node_t;

extern kalloc_node_t* kalloc_start;

static void wipeout_kalloc()
{
  memset(kalloc_start, 0, 1024);
  kalloc_start = NULL;
}

static struct TEST_STRUCT* new_test_struct()
{
  struct TEST_STRUCT* ret = kalloc(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 = kalloc(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, kalloc)
{
#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);

  char buf[1024];
  if (debug_kalloc_assert_consistency(buf, 1024)) {
    fprintf(stderr, "Consistency check failed. (%s:%d)\n", __FILE__, __LINE__);
    fprintf(stderr, buf);
    ASSERT_TRUE(false);
  }

  wipeout_kalloc();

  return 0;
}

struct UNEVEN_STRUCT {
  uint8_t arr[5];
};

struct UNEVEN_STRUCT* new_uneven_struct()
{
  struct UNEVEN_STRUCT* ret = kalloc(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_kalloc)
{
  if (kalloc_start) {
    wipeout_kalloc();
  }

  struct UNEVEN_STRUCT* test1 = new_uneven_struct();
  struct UNEVEN_STRUCT* test2 = new_uneven_struct();

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

  wipeout_kalloc();

  return 0;
}

TEST(memory, kalloc_free)
{
  if (kalloc_start) {
    wipeout_kalloc();
  }

  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();

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

  ASSERT_EQ((int)kalloc_start->size * 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();

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

  ASSERT_EQ((int)kalloc_start->size * 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();

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

  ASSERT_EQ((int)kalloc_start->size * 4, MAX_HEAP_SIZE - 4);

  wipeout_kalloc();

  return 0;
}

TEST(memory, kalloc_free_alloc2)
{
  if (kalloc_start) {
    wipeout_kalloc();
  }

  struct TEST_STRUCT2* test1 = new_test_struct2();
  struct TEST_STRUCT2* test2 = new_test_struct2();

  kfree(test1);

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

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

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

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

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

  return 0;
}

TEST(memory, relink_backref_after_free)
{
  if (kalloc_start) {
    wipeout_kalloc();
  }

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

  kfree(test2);
  kfree(test3);

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

  return 0;
}

TEST(memory, consistency_stress)
{
#define NRUNS 500
  if (kalloc_start) {
    wipeout_kalloc();
  }

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

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

    if (debug_kalloc_assert_consistency(buf, 1024)) {
      fprintf(
          stderr,
          "Consistency check failed. (At index=%d, %s:%d)\n",
          i,
          __FILE__,
          __LINE__);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }

    ASSERT_TRUE(allocd[i]);

    memset(allocd[i], 0xFF, nalloc);
    size_t idx = rand() % NRUNS;

    if (debug_kalloc_assert_consistency(buf, 1024)) {
      fprintf(
          stderr,
          "Consistency check failed. (At index=%d, %s:%d)\n",
          i,
          __FILE__,
          __LINE__);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }

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

    if (debug_kalloc_assert_consistency(buf, 1024)) {
      fprintf(
          stderr,
          "Consistency check failed. (At index=%d, %s:%d)\n",
          i,
          __FILE__,
          __LINE__);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }

    idx = rand() % NRUNS;
    kfree(allocd[idx]);
    allocd[idx] = NULL;

    if (debug_kalloc_assert_consistency(buf, 1024)) {
      fprintf(
          stderr,
          "Consistency check failed. (At index=%d, %s:%d)\n",
          i,
          __FILE__,
          __LINE__);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }
  }

  for (i = 0; i < NRUNS; ++i) {
    if (allocd[i]) {
      kfree(allocd[i]);
    }

    if (debug_kalloc_assert_consistency(buf, 1024)) {
      fprintf(
          stderr,
          "Consistency check failed. (At index=%d, %s:%d)\n",
          i,
          __FILE__,
          __LINE__);
      fprintf(stderr, buf);
      ASSERT_TRUE(false);
    }
  }
  ASSERT_EQ((int)kalloc_start->size * 4, MAX_HEAP_SIZE - 4);

  return 0;
}

TEST(memory, kalloc_free_alloc)
{
  if (kalloc_start) {
    wipeout_kalloc();
  }

  new_test_struct();
  struct TEST_STRUCT2* test2 = new_test_struct2();
  new_test_struct();
  struct TEST_STRUCT2* test4 = new_test_struct2();
  new_test_struct2();

  kfree(test4);

  struct TEST_STRUCT2* test6 = new_test_struct2();

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

  kfree(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;
}