#include <stdlib.h>

#include "shared/map.h"
#include "test_harness.h"

typedef const char* const_char_ptr;

#define integer_cmp_(a, b) (a - b)
MAP_DECL(int, const_char_ptr);
MAP_IMPL(int, const_char_ptr, integer_cmp_, null_dtor);

TEST(map, smoke)
{
  map_t(int, const_char_ptr)* map = map_new(int, const_char_ptr)();

  map_put(int, const_char_ptr)(map, 6, "string1");
  map_put(int, const_char_ptr)(map, 8, "string2");
  map_put(int, const_char_ptr)(map, 9, "string3");

  const char* str = *map_get(int, const_char_ptr)(map, 8);
  ASSERT_EQ_STR(str, "string2");

  str = *map_get(int, const_char_ptr)(map, 9);
  ASSERT_EQ_STR(str, "string3");
  ASSERT_EQ(map_get(int, const_char_ptr)(map, 20), NULL);

  map_free(int, const_char_ptr)(map);
  ASSERT_TRUE(debug_is_heap_empty());
  return 0;
}

typedef uint32_t* int_ptr;
MAP_DECL(int, int_ptr);
MAP_IMPL(int, int_ptr, integer_cmp_, kfree);

static inline void loose_map_print(
    avl_tree_node_t(map_entry_t(int, int_ptr))* node,
    int tab) {
  if (node) {
    loose_map_print(node->right, tab + 1);
  }
  for (int i = 0; i < tab; ++i) {
    printf(" ");
  }
  if (!node) {
    printf("(nil)\n");
    return;
  }
  printf("(%d => %d)\n", node->value.key, *node->value.value);
  loose_map_print(node->left, tab + 1);
}
static inline void map_debug_print(
    map_t(int, int_ptr) * map)
{
  printf("------\n");
  if (map->tree) {
    loose_map_print(map->tree->root, 0);
  }
  printf("------\n");
}
TEST(map, stress)
{
#define N 256
  uint32_t* arr[N];
  memset(arr, 0, sizeof(arr));
  map_t(int, int_ptr)* map = map_new(int, int_ptr)();

  for (int i = 0; i < N; ++i) {
    uint32_t* a = kalloc(sizeof(uint32_t));
    *a = rand();
    arr[i] = a;
  }

  for (int i = 0; i < N; ++i) {
    map_put(int, int_ptr)(map, i, arr[i]);
  }

  for (int i = 0; i < N; ++i) {
    uint32_t** ptr = map_get(int, int_ptr)(map, i);
    ASSERT_TRUE(ptr != NULL);
    ASSERT_TRUE(*ptr != NULL);
    ASSERT_EQ(**ptr, *arr[i]);
  }

  int erased[N];
  memset(erased, 0, sizeof(erased));

  for (int i = 0; i < (N / 2); ++i) {
    int_ptr out;
    int e = rand() % N;
    if (map_erase(int, int_ptr)(map, e, &out)) {
      ASSERT_EQ(*out, *arr[e]);
      arr[e] = NULL;
      kfree(out);
    }
    erased[e] = 1;
  }

  for (int i = 0; i < N; ++ i) {
    if (erased[i]) {
      continue;
    }
    uint32_t** ptr = map_get(int, int_ptr)(map, i);
    ASSERT_TRUE(ptr != NULL);
    ASSERT_TRUE(*ptr != NULL);
    ASSERT_EQ(**ptr, *arr[i]);
  }

  map_free(int, int_ptr)(map);
  ASSERT_TRUE(debug_is_heap_empty());

  return 0;
}