/**
 * @file HashTable.h
 * @author apio (cloudapio.eu)
 * @brief Table of values with best-case constant time lookup.
 *
 * @copyright Copyright (c) 2023, the Luna authors.
 *
 */

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

/**
 * @brief A table of values with best-case constant time lookup.
 *
 * @tparam T The type of values to store.
 */
template <typename T> class HashTable
{
    static constexpr usize GROW_RATE = 2;
    static constexpr usize GROW_FACTOR = 2;
    static constexpr usize INITIAL_CAPACITY = 32;

  public:
    /**
     * @brief Try to insert a value into the table.
     *
     * @param value The value to insert.
     * @return Result<bool> An error, true if the insertion succeeded, or false if the value already existed.
     */
    Result<bool> try_set(const T& value)
    {
        T copy { value };
        return try_set(move(copy));
    }

    /**
     * @brief Try to insert a value into the table.
     *
     * @param value The value to insert.
     * @return Result<bool> An error, true if the insertion succeeded, or false if the value already existed.
     */
    Result<bool> try_set(T&& value)
    {
        if (should_grow()) TRY(rehash(m_capacity ? m_capacity * GROW_FACTOR : INITIAL_CAPACITY));

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

        check(m_size < m_capacity);

        // NOTE: This should not end in an infinite loop if the table is full, since the table CANNOT be full
        // (should_grow() returns true before it does). The check above verifies this.
        while (true)
        {
            auto& bucket = m_buckets[index];
            if (bucket.has_value())
            {
                if (*bucket == value) return false;
                index++;
                if (index == m_capacity)
                    index = 0; // Wrap around to avoid overflowing, seems like I assumed it would do that
                               // automatically for some reason... (facepalm).
                continue;
            }
            bucket = { move(value) };
            m_size++;
            break;
        }

        return true;
    }

    /**
     * @brief Find a value inside the table, returning a pointer to it if found.
     *
     * @param value The value to compare against.
     * @return T* A pointer to the value inside the table if found, or nullptr if it was not found.
     */
    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++;
                if (i == m_capacity) i = 0;
            }
            else
                return nullptr;
        } while (i != index);

        return nullptr;
    }

    /**
     * @brief Remove a value from this table.
     *
     * @param value The value to remove.
     * @return true The value was successfully removed.
     * @return false The value did not exist.
     */
    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++;
                if (i == m_capacity) i = 0;
            }
            else
                return false;
        } while (i != index);

        return false;
    }

    /**
     * @brief Clear the table.
     */
    void clear()
    {
        if (m_capacity)
        {
            for (usize i = 0; i < m_capacity; i++) m_buckets[i].~Option();

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

    /**
     * @brief Return the number of values that can currently fit in the table.
     *
     * The number of actual entries will always be smaller than this (unless it is 0), since the table grows before
     * the number of entries reaches the capacity.
     *
     * @return usize The current capacity.
     */
    usize capacity() const
    {
        return m_capacity;
    }

    /**
     * @brief Return the number of values currently contained in the table.
     *
     * @return usize The number of values contained.
     */
    usize size() const
    {
        return m_size;
    }

    ~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 };
};