#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[16];
};

/* Copy of the node structure. */
// The sizes will count the number of WORDS allocated.
// Since there's a max size of 16k, only 12 bits will be
// needed for this.
typedef struct KALLOC_NODE {
  uint16_t size_words;
  uint16_t prev;
  /*
   * LSB is whether this block is used.
   *
   * Rest of the bits must equal (0xdeadbeee >> 1)
   */
  uint32_t used_and_canary;

  uint8_t mem[]; /* The memory to use. */
} PACKED 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) + sizeof(kalloc_node_t), 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, "%s", 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) + 8 + sizeof(kalloc_node_t), test2);

  wipeout_kalloc();

  return 0;
}

#define ASSERT_KALLOC_EMPTY()                                                                                \
  ASSERT_EQ(                                                                                                 \
      (void*)((uint8_t*)kalloc_start + kalloc_start->size_words * sizeof(uint32_t) + sizeof(kalloc_node_t)), \
      (void*)&HEAP_STOP)

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

  kalloc_init();

  ASSERT_KALLOC_EMPTY();

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

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

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

  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, "%s", 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, "%s", buf);
    ASSERT_TRUE(false);
  }

  return 0;
}

TEST(memory, alloc1)
{
  char buf[1024];

  kalloc(5);
  kalloc(1);

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

  return 0;
}

TEST(memory, alloc0)
{
  char buf[1024];

  kalloc(5);
  kalloc(0);

  if (debug_kalloc_assert_consistency(buf, 1024)) {
    fprintf(
        stderr,
        "Consistency check failed. (%s:%d)\n",
        __FILE__,
        __LINE__);
    fprintf(stderr, "%s", 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];
  void* tofree;

  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, nalloc=%ld)\n",
          i,
          __FILE__,
          __LINE__,
          nalloc);
      fprintf(stderr, "%s", 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, "%s", buf);
      ASSERT_TRUE(false);
    }

    tofree = allocd[idx];
    kfree(tofree);
    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, "%s", buf);
      ASSERT_TRUE(false);
    }

    idx = rand() % NRUNS;
    tofree = allocd[idx];
    kfree(tofree);
    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, "%s", 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, "%s", buf);
      ASSERT_TRUE(false);
    }
  }

  ASSERT_KALLOC_EMPTY();

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

  ASSERT_EQ(V(test7), V(test2));
  return 0;
}