/*
 * libfud
 * Copyright 2025 Dominick Allen
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUD_HASH_MAP_HPP
#define FUD_HASH_MAP_HPP

#include "fud_allocator.hpp"
#include "fud_hash.hpp"
#include "fud_hash_map_impl.hpp"
#include "fud_option.hpp"
#include "fud_result.hpp"
#include "fud_status.hpp"

#include <algorithm>

namespace fud {

constexpr double hashMapMaxLoadFactor = 0.6;

/** \brief An open-address hash map using quadratic probing.
 *
 * Templates on a Hash object to facilitate a variety of hashing function implementations. */
template <typename Key, typename Value, typename Hash>
class HashMap {
    static_assert(!std::is_same_v<Key, option_detail::NullOptionType>);
    static_assert(!std::is_same_v<Value, option_detail::NullOptionType>);

  public:
    struct KeyValuePair {
        Key m_key;
        Value m_value;
    };

    using Entry = Option<KeyValuePair>;
    using Node = detail::MapEntry<Key, Value>;
    static constexpr size_t NodeSize = sizeof(Node);
    static constexpr size_t Alignment = alignof(Node);

    constexpr HashMap() noexcept = default;

    /** \brief Construct a HashMap explicitly using the specified allocator. */
    constexpr explicit HashMap(Allocator& allocator) noexcept : m_allocator{&allocator}
    {
    }

    constexpr HashMap(const HashMap& rhs) = delete;

    constexpr HashMap(HashMap&& rhs) noexcept :
        m_allocator{rhs.m_allocator},
        m_data{rhs.m_data},
        m_count{rhs.m_count},
        m_buckets{rhs.m_buckets},
        m_seed{rhs.m_seed},
        m_hasher{std::move(rhs.m_hasher)}
    {
        rhs.m_allocator = nullptr;
        rhs.m_data = nullptr;
        rhs.m_count = 0;
        rhs.m_buckets = 0;
    }

    HashMap& operator=(const HashMap& rhs) = delete;

    HashMap& operator=(HashMap&& rhs) noexcept
    {
        if (&rhs == this) {
            return *this;
        }

        static_cast<void>(cleanup());
        m_allocator = rhs.m_allocator;
        m_data = rhs.m_data;
        m_count = rhs.m_count;
        m_buckets = rhs.m_buckets;
        m_seed = rhs.m_seed;
        m_hasher = std::move(rhs.m_hasher);

        rhs.m_allocator = nullptr;
        rhs.m_data = nullptr;
        rhs.m_count = 0;
        rhs.m_buckets = 0;

        return *this;
    }

    constexpr ~HashMap() noexcept
    {
        static_cast<void>(cleanup());
    }

    [[nodiscard]] size_t size() noexcept
    {
        return m_count;
    }

    [[nodiscard]] size_t capacity()
    {
        return m_buckets;
    }

    [[nodiscard]] bool contains(const Key& key) const
    {
        return lookup(key).hasValue();
    }

    [[nodiscard]] bool contains(Key&& key) const
    {
        return lookup(std::move(key)).hasValue();
    }

    FudStatus insert(const Key& key, const Value& value)
    {
        auto growStatus = checkGrow();
        if (growStatus != FudStatus::Success) {
            return growStatus;
        }

        auto hashIndexResult = findEmptyBucket(key, m_buckets, m_data);
        if (hashIndexResult.isError()) {
            return hashIndexResult.takeError();
        }
        auto hashIndex{hashIndexResult.takeOkay()};
        m_data[hashIndex].~Node();
        auto* ptr = new (m_data + hashIndex) Node{key, value};
        m_count++;
        return FudStatus::Success;
    }

    FudStatus insert(const Key& key, Value&& value)
    {
        auto growStatus = checkGrow();
        if (growStatus != FudStatus::Success) {
            return growStatus;
        }

        auto hashIndexResult = findEmptyBucket(key, m_buckets, m_data);
        if (hashIndexResult.isError()) {
            return hashIndexResult.takeError();
        }
        auto hashIndex{hashIndexResult.takeOkay()};
        m_data[hashIndex].~Node();
        auto* ptr = new (m_data + hashIndex) Node{key, std::move(value)};
        m_count++;
        return FudStatus::Success;
    }

    FudStatus insert(Key&& key, const Value& value)
    {
        auto growStatus = checkGrow();
        if (growStatus != FudStatus::Success) {
            return growStatus;
        }

        auto hashIndexResult = findEmptyBucket(key, m_buckets, m_data);
        if (hashIndexResult.isError()) {
            return hashIndexResult.takeError();
        }
        auto hashIndex{hashIndexResult.takeOkay()};
        m_data[hashIndex].~Node();
        auto* ptr = new (m_data + hashIndex) Node{std::move(key), value};
        m_count++;
        return FudStatus::Success;
    }

    FudStatus insert(Key&& key, Value&& value)
    {
        auto growStatus = checkGrow();
        if (growStatus != FudStatus::Success) {
            return growStatus;
        }

        auto hashIndexResult = findEmptyBucket(key, m_buckets, m_data);
        if (hashIndexResult.isError()) {
            return hashIndexResult.takeError();
        }
        auto hashIndex{hashIndexResult.takeOkay()};
        m_data[hashIndex].~Node();
        auto* ptr = new (m_data + hashIndex) Node{std::move(key), std::move(value)};
        m_count++;
        return FudStatus::Success;
    }

    FudStatus update(const Key& key, const Value& value)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return insert(key, value);
        }
        auto hashIndex{hashIndexOption.value()};
        m_data[hashIndex].value() = value;
        return FudStatus::Success;
    }

    FudStatus update(const Key& key, Value&& value)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return insert(key, std::move(value));
        }
        auto hashIndex{hashIndexOption.value()};
        m_data[hashIndex].value() = std::move(value);
        return FudStatus::Success;
    }

    FudStatus update(Key&& key, const Value& value)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return insert(std::move(key), value);
        }
        auto hashIndex{hashIndexOption.value()};
        m_data[hashIndex].value() = value;
        return FudStatus::Success;
    }

    FudStatus update(Key&& key, Value&& value)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return insert(std::move(key), std::move(value));
        }
        auto hashIndex{hashIndexOption.value()};
        m_data[hashIndex].value() = std::move(value);
        return FudStatus::Success;
    }

    FudStatus remove(const Key& key)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return FudStatus::NotFound;
        }

        m_data[hashIndexOption.value()].tombstone();

        return FudStatus::Success;
    }

    FudStatus remove(Key&& key)
    {
        auto hashIndexOption = lookup(std::move(key));
        if (hashIndexOption.isNone()) {
            return FudStatus::NotFound;
        }

        m_data[hashIndexOption.value()].tombstone();

        return FudStatus::Success;
    }

    Option<Value> extract(const Key& key)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        auto hashIndex{hashIndexOption.value()};
        Option<Value> value{std::move(m_data[hashIndex].takeValue())};
        m_data[hashIndex].tombstone();

        return value;
    }

    Option<Value> extract(Key&& key)
    {
        auto hashIndexOption = lookup(std::move(key));
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        auto hashIndex{hashIndexOption.value()};
        Option<Value> value{std::move(m_data[hashIndex].takeValue())};
        m_data[hashIndex].tombstone();

        return value;
    }

    Option<KeyValuePair> extractPair(const Key& key)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        auto hashIndex{hashIndexOption.value()};
        KeyValuePair kvPair{std::move(m_data[hashIndex].takeKey()), std::move(m_data[hashIndex].takeValue())};
        m_data[hashIndex].tombstone();

        return kvPair;
    }

    Option<KeyValuePair> extractPair(Key&& key)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        auto hashIndex{hashIndexOption.value()};
        KeyValuePair kvPair{std::move(key), std::move(m_data[hashIndex].takeValue())};
        m_data[hashIndex].tombstone();

        return kvPair;
    }

    Option<Value> get(const Key& key)
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        return m_data[hashIndexOption.value()].value();
    }

    Option<Value&> getRef(const Key& key) const
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        return m_data[hashIndexOption.value()].value();
    }

    Option<const Value&> getConstRef(const Key& key) const
    {
        auto hashIndexOption = lookup(key);
        if (hashIndexOption.isNone()) {
            return NullOpt;
        }

        return m_data[hashIndexOption.value()].value();
    }

    [[nodiscard]] bool empty() const;

    FudStatus clear()
    {
        if (m_allocator == nullptr || m_data == nullptr) {
            if (m_buckets > 0 || m_count > 0) {
                return FudStatus::ObjectInvalid;
            }
            return FudStatus::Success;
        }
        for (size_t index = 0; index < m_buckets; ++index) {
            m_data[index].~Node();
        }
        m_count = 0;
        return FudStatus::Success;
    }

    FudStatus reserve(size_t count)
    {
        if (count <= m_buckets) {
            return FudStatus::Success;
        }

        if (m_allocator == nullptr) {
            return FudStatus::ObjectInvalid;
        }

        if (count > SIZE_MAX / NodeSize) {
            return FudStatus::ArgumentInvalid;
        }

        size_t requestedSize = count * NodeSize;
        auto dataPtrResult = m_allocator->allocate(requestedSize, Alignment);
        if (dataPtrResult.isError()) {
            return dataPtrResult.takeError();
        }

        auto* dataPtr = std::bit_cast<Node*>(dataPtrResult.takeOkay());
        for (size_t index = 0; index < count; ++index) {
            const auto* ptr = new (dataPtr + index) Node();
            fudAssert(ptr != nullptr);
        }

        for (size_t index = 0; index < m_buckets; ++index) {
            if (m_data[index].hasValue()) {
                const auto& key = m_data[index].key();
                auto newHashIndexResult = findEmptyBucket(key, count, dataPtr);
                if (newHashIndexResult.isError()) {
                    m_allocator->deallocate(std::bit_cast<std::byte*>(dataPtr), requestedSize);
                    return FudStatus::Failure;
                }
                const auto newHashIndex{newHashIndexResult.takeOkay()};
                dataPtr[newHashIndex].~Node();
                const auto* ptr = new (dataPtr + newHashIndex) Node(std::move(m_data[index]));
                fudAssert(ptr != nullptr);
                m_data[index].~Node();
            }
        }

        auto status = FudStatus::Success;
        if (m_buckets > 0) {
            m_allocator->deallocate(std::bit_cast<std::byte*>(m_data), m_buckets * NodeSize);
        }

        m_data = dataPtr;
        m_buckets = count;

        return status;
    }

    Result<Option<Value>, FudStatus> exchange(const Key& key, const Value& value);

  protected:
    [[nodiscard]] double loadFactor() const
    {
        if (m_buckets == 0) {
            return hashMapMaxLoadFactor + 1.0;
        }
        return static_cast<double>(m_count) / static_cast<double>(m_buckets);
    }

    FudStatus checkGrow()
    {
        if (loadFactor() > hashMapMaxLoadFactor) {
            auto growStatus = grow();
            if (growStatus != FudStatus::Success) {
                return growStatus;
            }
        }
        return FudStatus::Success;
    }

    FudStatus grow()
    {
        size_t additional = m_buckets < 3 ? 3 : m_buckets / 2;
        constexpr auto maxSize = std::numeric_limits<size_t>::max();
        if (maxSize - additional * NodeSize < m_buckets * NodeSize) {
            additional = maxSize - m_buckets * NodeSize / 2;
        }
        while (additional > 0) {
            auto reserveStatus = reserve(additional + m_buckets);
            if (reserveStatus == FudStatus::Success) {
                break;
            }
            if (reserveStatus == FudStatus::AllocFailure) {
                additional /= 2;
            } else {
                return reserveStatus;
            }
        }

        if (loadFactor() >= hashMapMaxLoadFactor) {
            return FudStatus::AllocFailure;
        }

        return FudStatus::Success;
    }

    FudStatus cleanup() noexcept
    {
        auto status = clear();
        if (m_data != nullptr && m_allocator != nullptr) {
            m_allocator->deallocate(std::bit_cast<std::byte*>(m_data), m_buckets);
        }

        m_allocator = nullptr;
        m_data = nullptr;
        m_count = 0;
        m_buckets = 0;

        return status;
    }

    [[nodiscard]] static constexpr size_t calculateHashIndex(size_t hash, size_t probe)
    {
        return hash + ((probe + probe * probe) / 2);
    }

    Option<size_t> lookup(const Key& key) const
    {
        if (m_count == 0 || m_buckets == 0 || m_data == nullptr) {
            return NullOpt;
        }
        const auto hash = m_hasher(key, m_seed);
        const auto nearestPowerOf2 = detail::roundToNearest2(m_buckets);
        size_t probe = 0;
        size_t hashIndex{};
        auto collision = true;

        while (collision and probe < m_buckets) {
            hashIndex = calculateHashIndex(hash, probe) % nearestPowerOf2;
            probe++;
            if (hashIndex >= m_buckets) {
                continue;
            }
            collision = not m_data[hashIndex].isNone();
            if (m_data[hashIndex].isTombstone()) {
                continue;
            }
            if (collision and m_data[hashIndex].key() == key) {
                break;
            }
        }

        if (collision and hashIndex < m_buckets and m_data[hashIndex].key() == key) {
            return hashIndex;
        }

        return NullOpt;
    }

    Result<size_t, FudStatus> findEmptyBucket(const Key& key, size_t buckets, Node* data)
    {
        const auto hash = m_hasher(key, m_seed);
        const auto nearestPowerOf2 = detail::roundToNearest2(buckets);
        size_t hashIndex{};
        auto collision = true;
        size_t probe = 0;

        bool foundFirstSlot{false};
        size_t firstHashIndex{};
        while (collision and probe < buckets) {
            hashIndex = calculateHashIndex(hash, probe) % nearestPowerOf2;
            probe++;
            if (hashIndex >= buckets) {
                continue;
            }
            collision = not data[hashIndex].isNone();
            if (not foundFirstSlot and data[hashIndex].isTombstone()) {
                foundFirstSlot = true;
                firstHashIndex = hashIndex;
            }
            if (collision and data[hashIndex].key() == key) {
                break;
            }
        }

        if (collision and data[hashIndex].key() == key) {
            return Error{FudStatus::Exists};
        }

        if (foundFirstSlot) {
            return Okay{firstHashIndex};
        }

        if (collision) {
            return Error{FudStatus::Failure};
        }

        return Okay{hashIndex};
    }

    Allocator* m_allocator{&globalFudAllocator};
    Node* m_data{nullptr};
    size_t m_count{0};
    size_t m_buckets{0};
    size_t m_seed{0};
    Hash m_hasher{};
};

} // namespace fud

#endif