#include "Hash.h"
#include "Profiling.h"
#include <stdlib.h>
#include <string.h>

Hash::Hash() : m_bucket_count(0), m_used_buckets(0), m_keys(nullptr), m_values(nullptr)
{
    ZoneScoped;
}

Hash::Hash(uint32_t bucket_count) : m_bucket_count(0), m_used_buckets(0), m_keys(nullptr), m_values(nullptr)
{
    ZoneScoped;
    auto* mem = reinterpret_cast<uint64_t*>(malloc(sizeof(uint64_t) * 2 * bucket_count));
    if (!mem)
        return;
    m_bucket_count = bucket_count;
    m_keys = mem;
    m_values = mem + bucket_count;
    memset(mem, 0, sizeof(uint64_t) * 2 * bucket_count);
}

Hash::Hash(const Hash& copy)
    : m_bucket_count(copy.m_bucket_count),
      m_used_buckets(copy.m_used_buckets),
      m_keys(nullptr),
      m_values(nullptr)
{
    ZoneScoped;
    auto* mem = reinterpret_cast<uint64_t*>(malloc(sizeof(uint64_t) * 2 * m_bucket_count));
    if (!mem)
        return;
    m_keys = mem;
    m_values = mem + m_bucket_count;
    memcpy(m_keys, copy.m_keys, sizeof(uint64_t) * 2 * m_bucket_count);
}

Hash& Hash::operator=(const Hash& rhs)
{
    ZoneScoped;
    if (m_keys)
        free(m_keys);
    m_bucket_count = rhs.m_bucket_count;
    m_used_buckets = rhs.m_used_buckets;
    auto* mem = reinterpret_cast<uint64_t*>(malloc(sizeof(uint64_t) * 2 * m_bucket_count));
    if (!mem) {
        m_bucket_count = 0;
        m_keys = nullptr;
        m_values = nullptr;
        return *this;
    }
    m_keys = mem;
    m_values = mem + m_bucket_count;
    memcpy(m_keys, rhs.m_keys, sizeof(uint64_t) * 2 * m_bucket_count);
    return *this;
}

Hash::~Hash()
{
    ZoneScoped;
    free(m_keys);
    m_used_buckets = 0;
    m_bucket_count = 0;
    m_keys = nullptr;
    m_values = nullptr;
}

uint64_t Hash::lookup(uint64_t key, uint64_t default_value)
{
    ZoneScoped;
    if (m_bucket_count == 0)
        return default_value;
    uint64_t i = key % m_bucket_count;
    do {
        if (m_keys[i] == key)
            return m_values[i];
        else if (m_keys[i] == HASH_EMPTY_KEY)
            return default_value;
        i = (i + 1) % m_bucket_count;
    } while (i != key % m_bucket_count);
    return default_value;
}

bool Hash::insert(uint64_t key, uint64_t value)
{
    ZoneScoped;
    if (key == HASH_EMPTY_KEY || key == HASH_TOMBSTONE_KEY) {
        return false;
    }

    // Resize if necessary
    float usage = static_cast<float>(m_used_buckets) / static_cast<float>(m_bucket_count);
    if (m_bucket_count == 0 || usage > .5f) {
        uint32_t new_cap = (m_bucket_count > 0) ? m_bucket_count * 2 : 128;
        auto* new_mem = reinterpret_cast<uint64_t*>(malloc(sizeof(uint64_t) * new_cap * 2));
        if (!new_mem)
            return false;
        uint64_t* new_keys = new_mem;
        uint64_t* new_values = new_keys + new_cap;
        memset(new_keys, HASH_EMPTY_KEY, sizeof(uint64_t) * new_cap);

        // Re-insert all values
        for (uint32_t i = 0; i < m_bucket_count; ++i) {
            if (m_keys[i] != HASH_EMPTY_KEY) {
                uint64_t idx = m_keys[i] % new_cap;
                while (true) {
                    if (new_keys[idx] == HASH_EMPTY_KEY) {
                        new_keys[idx] = m_keys[i];
                        new_values[idx] = m_values[i];
                        break;
                    }
                    idx = (idx + 1) % new_cap;
                }
            }
        }
        free(m_keys);
        m_keys = new_keys;
        m_values = new_values;
        m_bucket_count = new_cap;
    }

    uint64_t i = key % m_bucket_count;
    do {
        if (m_keys[i] == HASH_EMPTY_KEY || m_keys[i] == HASH_TOMBSTONE_KEY) {
            m_keys[i] = key;
            m_values[i] = value;
            ++m_used_buckets;
            return true;
        }
        else if (m_keys[i] == key) {
            m_values[i] = value;
            return true;
        }
        i = (i + 1) % m_bucket_count;
    } while (i != (key % m_bucket_count));
    return false;
}

void Hash::remove(uint64_t key)
{
    ZoneScoped;
    if (m_bucket_count == 0)
        return;
    uint64_t i = key % m_bucket_count;
    do {
        if (m_keys[i] == key) {
            m_keys[i] = HASH_TOMBSTONE_KEY;
        }
        else if (m_keys[i] == HASH_EMPTY_KEY) {
            break;
        }
    } while (i != key % m_bucket_count);
}
