#pragma once

#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <optional>
#include <stdexcept>
#include <utility>
#include <vector>

namespace uc::detail {

// Flat open-addressed value -> duplicate-chain metadata.  The actual duplicate
// links live in AdaptiveSequence's dense ID table, so one bucket is paid per
// distinct value rather than one node allocation per element.
template <class Key, class Hash = std::hash<Key>, class Equal = std::equal_to<Key>>
class FlatDuplicateIndex {
public:
    using id_type = std::uint32_t;
    static constexpr id_type invalid_id = std::numeric_limits<id_type>::max();

    struct Entry {
        id_type head = invalid_id;
        std::uint32_t count = 0;
    };

    FlatDuplicateIndex() { rehash(16); }

    [[nodiscard]] const Entry* find(const Key& key) const {
        const auto position = find_existing(key);
        return position == npos ? nullptr : &buckets_[position].entry;
    }

    [[nodiscard]] Entry* find(const Key& key) {
        const auto position = find_existing(key);
        return position == npos ? nullptr : &buckets_[position].entry;
    }

    Entry& ensure(const Key& key) {
        maybe_grow();
        const auto [position, found] = find_insert_position(key);
        auto& bucket = buckets_[position];
        if (!found) {
            if (bucket.state == State::tombstone) {
                --tombstones_;
            }
            bucket.key.emplace(key);
            bucket.entry = {};
            bucket.state = State::occupied;
            ++size_;
        }
        return bucket.entry;
    }

    bool erase_key(const Key& key) {
        const auto position = find_existing(key);
        if (position == npos) {
            return false;
        }
        auto& bucket = buckets_[position];
        bucket.key.reset();
        bucket.entry = {};
        bucket.state = State::tombstone;
        --size_;
        ++tombstones_;
        if (tombstones_ > buckets_.size() / 4) {
            rehash(buckets_.size());
        }
        return true;
    }

    void clear() {
        buckets_.clear();
        size_ = 0;
        tombstones_ = 0;
        rehash(16);
    }

    [[nodiscard]] std::size_t distinct_values() const noexcept { return size_; }
    [[nodiscard]] std::size_t bucket_count() const noexcept { return buckets_.size(); }
    [[nodiscard]] std::size_t allocated_bytes() const noexcept {
        return buckets_.capacity() * sizeof(Bucket);
    }

private:
    enum class State : std::uint8_t { empty, occupied, tombstone };

    struct Bucket {
        std::optional<Key> key;
        Entry entry;
        State state = State::empty;
    };

    static constexpr std::size_t npos = std::numeric_limits<std::size_t>::max();

    [[nodiscard]] std::size_t mask() const noexcept { return buckets_.size() - 1; }

    [[nodiscard]] std::size_t find_existing(const Key& key) const {
        if (buckets_.empty()) {
            return npos;
        }
        auto position = hasher_(key) & mask();
        for (std::size_t probe = 0; probe < buckets_.size(); ++probe) {
            const auto& bucket = buckets_[position];
            if (bucket.state == State::empty) {
                return npos;
            }
            if (bucket.state == State::occupied && equal_(*bucket.key, key)) {
                return position;
            }
            position = (position + 1) & mask();
        }
        return npos;
    }

    [[nodiscard]] std::pair<std::size_t, bool>
    find_insert_position(const Key& key) const {
        auto position = hasher_(key) & mask();
        auto first_tombstone = npos;
        for (std::size_t probe = 0; probe < buckets_.size(); ++probe) {
            const auto& bucket = buckets_[position];
            if (bucket.state == State::empty) {
                return {first_tombstone == npos ? position : first_tombstone, false};
            }
            if (bucket.state == State::tombstone) {
                if (first_tombstone == npos) {
                    first_tombstone = position;
                }
            } else if (equal_(*bucket.key, key)) {
                return {position, true};
            }
            position = (position + 1) & mask();
        }
        if (first_tombstone != npos) {
            return {first_tombstone, false};
        }
        throw std::length_error("flat hash index is full");
    }

    void maybe_grow() {
        if ((size_ + tombstones_ + 1) * 100 >= buckets_.size() * 82) {
            rehash(buckets_.size() * 2);
        }
    }

    void rehash(std::size_t requested_capacity) {
        std::size_t capacity = 16;
        while (capacity < requested_capacity) {
            capacity *= 2;
        }
        auto old = std::move(buckets_);
        buckets_.assign(capacity, Bucket{});
        size_ = 0;
        tombstones_ = 0;
        for (auto& bucket : old) {
            if (bucket.state != State::occupied) {
                continue;
            }
            const auto [position, found] = find_insert_position(*bucket.key);
            (void)found;
            auto& target = buckets_[position];
            target.key.emplace(std::move(*bucket.key));
            target.entry = bucket.entry;
            target.state = State::occupied;
            ++size_;
        }
    }

    Hash hasher_{};
    Equal equal_{};
    std::vector<Bucket> buckets_;
    std::size_t size_ = 0;
    std::size_t tombstones_ = 0;
};

} // namespace uc::detail