#pragma once
#include <luna/Hash.h>
#include <luna/Heap.h>
#include <luna/Option.h>

template <typename T> class HashTable
{
    static constexpr usize GROW_RATE = 2;
    static constexpr usize GROW_FACTOR = 16;

  public:
    Result<bool> try_set(const T& value)
    {
        T copy { value };
        return try_set(move(copy));
    }

    Result<bool> try_set(T&& value)
    {
        if (should_grow()) TRY(rehash(m_capacity + GROW_FACTOR));

        u64 index = hash(value, m_salt) % m_capacity;

        while (true)
        {
            auto& bucket = m_buckets[index];
            if (bucket.has_value())
            {
                if (*bucket == value) return false;
                index++;
                continue;
            }
            bucket = { move(value) };
            m_size++;
            break;
        }

        return true;
    }

    T* try_find(const T& value)
    {
        if (!m_size) return nullptr;

        check(m_capacity);

        const u64 index = hash(value, m_salt) % m_capacity;
        usize i = index;

        do {
            auto& bucket = m_buckets[i];
            if (bucket.has_value())
            {
                if (*bucket == value) return bucket.value_ptr();
                i++;
            }
            else
                return nullptr;
        } while (i != index);

        return nullptr;
    }

    bool try_remove(const T& value)
    {
        if (!m_size) return false;

        check(m_capacity);

        const u64 index = hash(value, m_salt) % m_capacity;
        usize i = index;

        do {
            auto& bucket = m_buckets[i];
            if (bucket.has_value())
            {
                if (*bucket == value)
                {
                    bucket = {};
                    m_size--;
                    if (i != index) rehash(m_capacity);
                    return true;
                }
                i++;
            }
            else
                return false;
        } while (i != index);

        return false;
    }

    void clear()
    {
        for (usize i = 0; i < m_capacity; i++) m_buckets[i].~Option();

        free_impl(m_buckets);
        m_capacity = m_size = 0;
    }

    ~HashTable()
    {
        clear();
    }

  private:
    bool should_grow()
    {
        return (m_capacity == 0) || ((m_size * GROW_RATE) >= m_capacity);
    }

    Result<void> rehash(usize new_capacity)
    {
        HashTable<T> new_table;
        TRY(new_table.initialize(new_capacity));

        if (m_capacity != 0)
        {
            for (usize i = 0; i < m_capacity; i++)
            {
                auto& opt = m_buckets[i];
                if (opt.has_value())
                {
                    auto value = opt.release_value();
                    TRY(new_table.try_set(move(value)));
                }
            }
        }

        swap(this, &new_table);

        return {};
    }

    Result<void> initialize(usize initial_capacity)
    {
        check(m_buckets == nullptr);
        m_capacity = initial_capacity;
        m_buckets = (Option<T>*)TRY(calloc_impl(initial_capacity, sizeof(Option<T>), false));
        return {};
    }

    Option<T>* m_buckets { nullptr };
    usize m_capacity { 0 };
    usize m_size { 0 };
    // FIXME: Randomize this to protect against hash table attacks.
    u64 m_salt { 0 };
};