#define HM_IMPLEMENTATION
#include "hm.h"

#include <stdio.h>

static int failures = 0;

#define CHECK(cond)                                                    \
    do {                                                               \
        if (!(cond)) {                                                 \
            fprintf(stderr, "%s:%d: CHECK failed: %s\n",               \
                    __FILE__, __LINE__, #cond);                        \
            ++failures;                                                \
        }                                                              \
    } while (0)

// Forces every key into the same chain
static uint64_t const_hash(const void *buffer, size_t size)
{
    (void)buffer;
    (void)size;
    return 0;
}

static void test_int_keys(void)
{
    HashMap hm = hm_create(sizeof(int), sizeof(int));
    const int n = 10000;

    for (int i = 0; i < n; ++i) {
        int v = i * 2;
        CHECK(hm_put(&hm, &i, &v));
    }
    CHECK(hm_count(&hm) == (size_t)n);

    for (int i = 0; i < n; ++i) {
        int *v = hm_of(int, &hm, &i);
        CHECK(v != NULL && *v == i * 2);
    }

    int missing = n;
    CHECK(hm_get(&hm, &missing) == NULL);

    for (int i = 0; i < n; i += 2) {
        CHECK(hm_remove(&hm, &i));
    }
    CHECK(hm_count(&hm) == (size_t)n / 2);
    for (int i = 0; i < n; ++i) {
        CHECK((hm_get(&hm, &i) != NULL) == (i % 2 == 1));
    }

    hm_free(&hm);
}

static void test_string_value_resize_in_chain(void)
{
    HashMap hm = hm_create_ex(0, 0, const_hash);
    char big[4096];
    memset(big, 'y', sizeof(big) - 1);
    big[sizeof(big) - 1] = '\0';

    // Chain: c -> b -> a
    CHECK(hm_put(&hm, "a", "x"));
    CHECK(hm_put(&hm, "b", "x"));
    CHECK(hm_put(&hm, "c", "x"));

    // Grow every bucket so realloc moves them, then shrink them back
    CHECK(hm_put(&hm, "a", big));
    CHECK(hm_put(&hm, "b", big));
    CHECK(hm_put(&hm, "c", big));
    CHECK(strcmp((char*)hm_get(&hm, "b"), big) == 0);
    CHECK(hm_put(&hm, "b", "short"));
    CHECK(strcmp((char*)hm_get(&hm, "b"), "short") == 0);
    CHECK(hm_count(&hm) == 3);

    // Removing neighbors walks the prev/next pointers fixed up above
    CHECK(hm_remove(&hm, "a"));
    CHECK(hm_remove(&hm, "c"));
    CHECK(hm_remove(&hm, "b"));
    CHECK(hm_count(&hm) == 0);

    hm_free(&hm);
}

static void test_alignment(void)
{
    // Variable length string keys followed by doubles
    HashMap hm = hm_create(0, sizeof(double));
    const char *keys[] = { "a", "ab", "abc", "abcd", "abcde", "abcdefg" };
    for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); ++i) {
        double d = (double)i + 0.5;
        CHECK(hm_put(&hm, keys[i], &d));
        double *p = hm_of(double, &hm, keys[i]);
        CHECK(p != NULL && (uintptr_t)p % alignof(double) == 0);
        CHECK(p != NULL && *p == d);
    }
    hm_free(&hm);

    // Odd-sized struct keys followed by struct values
    typedef struct { int a, b, c; } Key;
    typedef struct { double x; long long y; } Value;
    hm = hm_create(sizeof(Key), sizeof(Value));
    for (int i = 0; i < 100; ++i) {
        Key k;
        memset(&k, 0, sizeof(k));
        k.a = i;
        Value v = { (double)i, i };
        CHECK(hm_put(&hm, &k, &v));
        CHECK((uintptr_t)hm_get(&hm, &k) % alignof(Value) == 0);
        CHECK((uintptr_t)hm_get(&hm, &k) % alignof(double) == 0);
        CHECK(hm_of(Value, &hm, &k)->y == i);
    }
    HashMapIterator it = hm_iterate(&hm);
    while (hm_next(&it) != NULL) {
        CHECK((uintptr_t)hm_key(&it) % alignof(int) == 0);
    }
    hm_free(&hm);
}

static void test_zero_initialized(void)
{
    HashMap hm = {0};

    CHECK(hm_get(&hm, "key") == NULL);
    CHECK(!hm_remove(&hm, "key"));
    CHECK(hm_put(&hm, "key", "value"));
    CHECK(strcmp((char*)hm_get(&hm, "key"), "value") == 0);

    hm_free(&hm);
}

static void test_reuse_after_free(void)
{
    HashMap hm = hm_create(sizeof(int), sizeof(int));
    int k = 1, v = 2;

    CHECK(hm_put(&hm, &k, &v));
    hm_free(&hm);
    CHECK(hm_count(&hm) == 0);
    CHECK(hm_get(&hm, &k) == NULL);

    // Configuration survives hm_free()
    v = 3;
    CHECK(hm_put(&hm, &k, &v));
    CHECK(*hm_of(int, &hm, &k) == 3);

    hm_free(&hm);
    hm_free(&hm);
}

static void test_iteration(void)
{
    HashMap hm = hm_create(sizeof(int), sizeof(int));
    const int n = 1000;
    static bool seen[1000];

    HashMapIterator it = hm_iterate(&hm);
    CHECK(hm_next(&it) == NULL);

    for (int i = 0; i < n; ++i) {
        CHECK(hm_put(&hm, &i, &i));
    }

    size_t visited = 0;
    it = hm_iterate(&hm);
    while (hm_next(&it) != NULL) {
        int key = *(const int*)hm_key(&it);
        CHECK(key >= 0 && key < n && !seen[key]);
        CHECK(*(int*)hm_value(&it) == key);
        seen[key] = true;
        ++visited;
    }
    CHECK(visited == (size_t)n);
    CHECK(hm_key(&it) == NULL && hm_value(&it) == NULL);

    hm_free(&hm);
}

static void test_update_does_not_rehash(void)
{
    HashMap hm = hm_create(sizeof(int), sizeof(int));
    const int n = (int)(HM_INITIAL_CAPACITY * HM_MAX_LOAD_FACTOR);

    for (int i = 0; i < n; ++i) {
        CHECK(hm_put(&hm, &i, &i));
    }
    CHECK(hm.capacity == HM_INITIAL_CAPACITY);

    int k = 0, v = 42;
    CHECK(hm_put(&hm, &k, &v));
    CHECK(hm.capacity == HM_INITIAL_CAPACITY);
    CHECK(*hm_of(int, &hm, &k) == 42);

    hm_free(&hm);
}

static void test_custom_hash(void)
{
    HashMap hm = hm_create_ex(0, sizeof(int), const_hash);
    char key[16];

    for (int i = 0; i < 500; ++i) {
        snprintf(key, sizeof(key), "k%d", i);
        CHECK(hm_put(&hm, key, &i));
    }
    for (int i = 0; i < 500; ++i) {
        snprintf(key, sizeof(key), "k%d", i);
        CHECK(*hm_of(int, &hm, key) == i);
    }
    hm_free(&hm);

    hm = hm_create_ex(0, 0, NULL);
    CHECK(hm.hash_func == HM_DEFAULT_HASH_FUNCTION);
    hm_free(&hm);
}

static size_t last_hash_size;

// Records the size it receives
static uint64_t size_hash(const void *buffer, size_t size)
{
    last_hash_size = size;
    return hm__fnv1a(buffer, size);
}

static void test_hash_size(void)
{
    // NULL terminated keys are hashed without their terminator
    HashMap hm = hm_create_ex(0, 0, size_hash);
    CHECK(hm_put(&hm, "abc", "v"));
    CHECK(last_hash_size == 3);
    CHECK(hm_get(&hm, "abc") != NULL);
    CHECK(last_hash_size == 3);
    CHECK(hm_put(&hm, "", "empty"));
    CHECK(last_hash_size == 0);
    CHECK(strcmp((char*)hm_get(&hm, ""), "empty") == 0);
    CHECK(hm_remove(&hm, ""));
    CHECK(last_hash_size == 0);
    hm_free(&hm);

    // Fixed-size keys are hashed with their full size
    hm = hm_create_ex(sizeof(int), sizeof(int), size_hash);
    int k = 7;
    CHECK(hm_put(&hm, &k, &k));
    CHECK(last_hash_size == sizeof(int));
    hm_free(&hm);
}

int main(void)
{
    test_int_keys();
    test_string_value_resize_in_chain();
    test_alignment();
    test_zero_initialized();
    test_reuse_after_free();
    test_iteration();
    test_update_does_not_rehash();
    test_custom_hash();
    test_hash_size();

    if (failures > 0) {
        fprintf(stderr, "%d check(s) failed\n", failures);
        return 1;
    }

    printf("All tests passed\n");
    return 0;
}