// this will certainly fail if NDEBUG is defined
#undef NDEBUG
#include <assert.h>
#include <ht.h>
#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

void *counting_malloc(size_t size, void *count_ptr) {
    int64_t *count = count_ptr;
    void *ptr = malloc(size);
    if (!ptr) {
        return NULL;
    }
    ++*count;
    return ptr;
}

void counting_free(void *ptr, void *count_ptr) {
    int64_t *count = count_ptr;
    if (ptr) {
        --*count;
    }
    free(ptr);
}

// this is void so it is standards compliant to pass this as a callback
void *counting_strdup(const void *str, void *count_ptr) {
    size_t len = strlen(str) + 1;
    char *buf = counting_malloc(len, count_ptr);
    if (!buf) {
        return NULL;
    }
    memcpy(buf, str, len);
    return buf;
}

struct count_and_strs {
    char **strs;
    int count;
};

bool foreach_test_callback(void *key, void *value, void *user_data) {
    struct count_and_strs *data = user_data;
    char *index_str = (char *) key + sizeof("str") - 1;
    int num = atoi(index_str) - 1;
    assert(strcmp(key, data->strs[num]) == 0);
    assert(strcmp(value, data->strs[num]) == 0);
    return ++data->count >= 50;
}

bool foreach_remove_test_callback(void *key, void *value, void *user_data) {
    bool *flag = user_data;
    return *flag = !*flag;
}

struct alloc_count_and_flag {
    bool flag;
    int64_t *count;
};

bool foreach_steal_test_callback(void *key, void *value, void *user_data) {
    struct alloc_count_and_flag *data = user_data;
    if (data->flag) {
        counting_free(key, data->count);
        counting_free(value, data->count);
    }
    data->flag = !data->flag;
    return !data->flag;
}

int main() {
    int64_t allocation_count = 0;
    const HTTableFunctions ONE_STR_FNS = {
        .hash = ht_string_hash_callback,
        .equal = ht_string_equal_callback,
        .destroy_key = counting_free,
        .destroy_value = NULL,
        .user_data = &allocation_count,
    };
    const HTTableFunctions TWO_STR_FNS = {
        .hash = ht_string_hash_callback,
        .equal = ht_string_equal_callback,
        .destroy_key = counting_free,
        .destroy_value = counting_free,
        .user_data = &allocation_count,
    };
    const HTAllocator ALLOCATOR = {
        .malloc = counting_malloc,
        .free = counting_free,
        .user_data = &allocation_count,
    };
    const HTTableFunctions INT_FNS = {
        .hash = ht_intptr_hash_callback,
        .equal = ht_intptr_equal_callback,
        .destroy_key = NULL,
        .destroy_value = NULL,
        .user_data = NULL,
    };
#define malloc(s) counting_malloc((s), &allocation_count)
#define free(p)   counting_free((p), &allocation_count)
#define strdup(s) counting_strdup((s), &allocation_count)

    HTTable *t = ht_new(&ONE_STR_FNS, &ALLOCATOR, NULL);
    assert(ht_count(t) == 0);

    assert(ht_insert(t, strdup("str1"), HT_STUFF(1)));
    assert(ht_has(t, "str1"));
    assert(ht_count(t) == 1);
    assert(ht_get(t, "str1") == HT_STUFF(1));

    assert(ht_insert(t, strdup("str2"), HT_STUFF(2)));
    assert(ht_count(t) == 2);
    assert(ht_get(t, "str2") == HT_STUFF(2));

    assert(ht_has(t, "str1"));
    assert(ht_has(t, "str2"));
    assert(!ht_has(t, "str3"));
    assert(ht_get(t, "str3") == NULL);

    assert(ht_remove(t, "str2"));
    assert(ht_count(t) == 1);
    assert(ht_remove(t, "str2"));
    assert(!ht_has(t, "str2"));
    assert(ht_get(t, "str2") == NULL);
    assert(ht_count(t) == 1);

    ht_free(t);

    t = ht_new(&ONE_STR_FNS, &ALLOCATOR, NULL);
    HTTable *t2 = ht_new(&ONE_STR_FNS, &ALLOCATOR, NULL);

    assert(ht_insert(t, strdup("1"), HT_STUFF(1)));
    assert(ht_insert(t, strdup("2"), HT_STUFF(2)));
    assert(ht_insert(t2, strdup("2"), HT_STUFF(4)));
    assert(ht_insert(t2, strdup("3"), HT_STUFF(3)));
    assert(ht_count(t) == 2);
    assert(ht_count(t2) == 2);

    assert(ht_steal_from(t, t2));
    assert(ht_count(t) == 3);
    assert(ht_count(t2) == 0);
    assert(ht_get(t, "1") == HT_STUFF(1));
    assert(ht_get(t, "2") == HT_STUFF(4));
    assert(ht_get(t, "3") == HT_STUFF(3));

    void *found_key;
    assert(ht_steal(t, "2", &found_key, NULL));
    assert(strcmp(found_key, "2") == 0);
    free(found_key);

    void *found_value;
    assert(ht_steal(t, "1", NULL, &found_value));
    assert(found_value == HT_STUFF(1));

    assert(ht_steal(t, "3", &found_key, &found_value));
    assert(strcmp(found_key, "3") == 0);
    assert(found_value == HT_STUFF(3));
    free(found_key);

    ht_free(t);
    ht_free(t2);

    t = ht_new(&INT_FNS, &ALLOCATOR, NULL);
    assert(ht_count(t) == 0);

    for (size_t i = 0; i < 500; ++i) {
        assert(ht_insert(t, HT_STUFF(i), HT_STUFF(-i)));
        assert(ht_count(t) == i + 1);
        assert(ht_has(t, HT_STUFF(i)));
        assert(ht_get(t, HT_STUFF(i)) == HT_STUFF(-i));
        assert(!ht_has(t, HT_STUFF(i + 1)));
    }

    for (size_t i = 0; i < 500; ++i) {
        assert(ht_has(t, HT_STUFF(i)));
    }

    for (size_t i = 0; i < 500; ++i) {
        assert(ht_remove(t, HT_STUFF(i)));
        assert(ht_count(t) == 499 - i);
        assert(!ht_has(t, HT_STUFF(i)));
        assert(ht_get(t, HT_STUFF(i)) == 0);
        if (i < 499) {
            assert(ht_has(t, HT_STUFF(i + 1)));
        }
    }

    ht_free(t);

    t = ht_new(&TWO_STR_FNS, &ALLOCATOR, NULL);

    char *STRS_COPY1[100];
    char *STRS_COPY2[100];
    char *STRS_COPY3[100];
    const size_t MAX_STR_LEN = sizeof("str100");

    for (size_t i = 0; i < 100; ++i) {
        STRS_COPY1[i] = malloc(MAX_STR_LEN);
        snprintf(STRS_COPY1[i], MAX_STR_LEN, "str%zu", i + 1);
        STRS_COPY2[i] = strdup(STRS_COPY1[i]);
        STRS_COPY3[i] = strdup(STRS_COPY1[i]);
        assert(ht_insert(t, STRS_COPY1[i], STRS_COPY2[i]));
    }

    assert(ht_count(t) == 100);

    for (size_t i = 0; i < 100; ++i) {
        void *found_key;
        assert(ht_get_extended(t, STRS_COPY3[i], &found_key) == STRS_COPY2[i]);
        assert(found_key == STRS_COPY1[i]);
    }

    t2 = ht_copy(t, counting_strdup, counting_strdup, &allocation_count);

    assert(ht_count(t) == 100);
    assert(ht_count(t2) == 100);
    assert(ht_equal(t, t2));

    for (size_t i = 0; i < 100; ++i) {
        assert(ht_remove(t, STRS_COPY3[i]));
        if (i < 5) {
            assert(!ht_equal(t, t2));
        }
    }

    assert(ht_count(t) == 0);
    assert(ht_count(t2) == 100);
    assert(!ht_equal(t, t2));

    ht_free(t);
    assert(ht_count(t2) == 100);

    for (size_t i = 0; i < 100; ++i) {
        void *found_key;
        assert(ht_has_extended(t2, STRS_COPY3[i], &found_key));
        assert(STRS_COPY3[i] != found_key);
    }

    void **keys1, **keys2;
    void **values1, **values2;
    assert((keys1 = ht_get_keys(t2)));
    assert((values1 = ht_get_values(t2)));
    assert(ht_get_keys_and_values(t2, &keys2, &values2));

    for (size_t i = 0; i < 100; ++i) {
        char *index_str = (char *) keys1[i] + sizeof("str") - 1;
        int num = atoi(index_str) - 1;
        assert(strcmp(keys1[i], keys2[i]) == 0);
        assert(strcmp(keys1[i], STRS_COPY3[num]) == 0);
        assert(strcmp(values1[i], values1[i]) == 0);
        assert(strcmp(values1[i], STRS_COPY3[num]) == 0);
    }

    free(keys1);
    free(keys2);
    free(values1);
    free(values2);

    struct count_and_strs cas = {
        .strs = STRS_COPY3,
        .count = 0,
    };
    ht_foreach(t2, foreach_test_callback, &cas);
    assert(cas.count == 50);

    assert(ht_count(t2) == 100);
    bool flag = NULL;
    assert(ht_foreach_remove(t2, foreach_remove_test_callback, &flag));
    assert(ht_count(t2) == 50);

    struct alloc_count_and_flag acaf = {
        .flag = false,
        .count = &allocation_count,
    };
    assert(ht_foreach_steal(t2, foreach_steal_test_callback, &acaf));
    assert(ht_count(t2) == 25);

    t = ht_copy(t2, counting_strdup, counting_strdup, &allocation_count);

    assert(ht_clear(t2));
    assert(ht_count(t2) == 0);
    ht_free(t2);

    for (size_t i = 0; i < 100; ++i) {
        free(STRS_COPY3[i]);
    }

    assert(ht_count(t) == 25);
    void **keys, **values;
    ht_steal_all(t, &keys, &values);
    assert(ht_count(t) == 0);

    for (size_t i = 0; i < 25; ++i) {
        assert(strcmp(keys[i], values[i]) == 0);
        assert(keys[i] != values[i]);
        free(keys[i]);
        free(values[i]);
    }
    free(keys);
    free(values);

    ht_free(t);

    assert(allocation_count == 0);
    return 0;
}