#pragma once

#include <cstddef>
#include <cstdint>
#include <functional>
#include <limits>
#include <optional>
#include <stdexcept>
#include <type_traits>
#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) {
        // A moved-from std::vector is allowed to be empty while the scalar
        // counters retain their old values.  Rebuild lazily so a moved-from
        // index remains reusable and mask() is never evaluated for zero
        // buckets.
        if (buckets_.empty()) {
            rehash(16);
        }

        auto [position, found] = find_insert_position(key);
        if (position == npos) {
            // Normal growth is performed by reserve_for_elements() at a
            // logical container-capacity change.  Reaching a completely full
            // table means that invariant was unavailable (most notably after
            // reusing a moved-from standalone index), so grow only as a
            // recovery path rather than as a load-factor policy.
            if (buckets_.size() > std::numeric_limits<std::size_t>::max() / 2) {
                throw std::length_error("flat hash index capacity overflow");
            }
            rehash(buckets_.size() * 2);
            const auto retry = find_insert_position(key);
            position = retry.first;
            found = retry.second;
            if (position == npos) {
                throw std::length_error("flat hash index is full");
            }
        }

        auto& bucket = buckets_[position];
        if (!found) {
            bucket.key.emplace(key);
            if (bucket.state == State::tombstone) {
                --tombstones_;
            }
            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;
        }
        if constexpr (std::is_nothrow_invocable_v<const Hash&, const Key&>
                      && std::is_nothrow_move_assignable_v<Bucket>) {
            erase_and_compact(position);
        } else {
            // Computing hashes or moving a key can throw.  Rebuild a complete
            // replacement without the erased key, then commit with swap, so a
            // failed erase never opens a hole in the live probe chain.
            erase_transactional(position);
        }
        --size_;
        return true;
    }

    void clear() {
        // Allocate before committing so a failed reset leaves the live table
        // and its duplicate metadata untouched.
        std::vector<Bucket> replacement(16);
        buckets_.swap(replacement);
        size_ = 0;
        tombstones_ = 0;
    }

    // Used by the owning container after moving the live table elsewhere.
    // Keeping zero buckets avoids an allocation in its noexcept move path;
    // ensure()/reserve_for_elements() restore the minimum table lazily.
    void reset_after_move() noexcept {
        buckets_.clear();
        size_ = 0;
        tombstones_ = 0;
    }

    [[nodiscard]] std::size_t distinct_values() const noexcept { return size_; }
    [[nodiscard]] std::size_t bucket_count() const noexcept { return buckets_.size(); }
    [[nodiscard]] static std::size_t buckets_for_elements(
        std::size_t element_capacity) {
        // At most one distinct key per element.  Keeping the worst-case load
        // below 70% guarantees that ensure() cannot trigger an independent
        // rehash between container capacity changes.
        const auto required = element_capacity > (std::numeric_limits<std::size_t>::max() - 6) / 10
            ? std::numeric_limits<std::size_t>::max()
            : (element_capacity * 10 + 6) / 7;
        std::size_t buckets = 16;
        while (buckets < required) {
            if (buckets > std::numeric_limits<std::size_t>::max() / 2) {
                throw std::length_error("flat hash index capacity overflow");
            }
            buckets *= 2;
        }
        return buckets;
    }
    void reserve_for_elements(std::size_t element_capacity) {
        const auto requested = buckets_for_elements(element_capacity);
        if (requested != buckets_.size() || tombstones_ != 0) {
            rehash(requested);
        }
    }
    [[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 {
        return find_insert_position_in(buckets_, key);
    }

    [[nodiscard]] std::pair<std::size_t, bool>
    find_insert_position_in(const std::vector<Bucket>& buckets,
                            const Key& key) const {
        if (buckets.empty()) {
            return {npos, false};
        }
        const auto table_mask = buckets.size() - 1;
        auto position = hasher_(key) & table_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) & table_mask;
        }
        if (first_tombstone != npos) {
            return {first_tombstone, false};
        }
        return {npos, false};
    }

    void erase_and_compact(std::size_t position) {
        auto hole = position;
        buckets_[hole] = Bucket{};
        auto current = (hole + 1) & mask();
        while (buckets_[current].state != State::empty) {
            auto& bucket = buckets_[current];
            if (bucket.state == State::occupied) {
                const auto home = hasher_(*bucket.key) & mask();
                const auto current_distance = (current - home) & mask();
                const auto hole_distance = (hole - home) & mask();
                if (hole_distance < current_distance) {
                    buckets_[hole] = std::move(bucket);
                    bucket = Bucket{};
                    hole = current;
                }
            }
            current = (current + 1) & mask();
        }
    }

    void erase_transactional(std::size_t erased_position) {
        std::vector<Bucket> replacement(buckets_.size());
        for (std::size_t source = 0; source < buckets_.size(); ++source) {
            const auto& bucket = buckets_[source];
            if (source == erased_position || bucket.state != State::occupied) {
                continue;
            }
            const auto [position, found] =
                find_insert_position_in(replacement, *bucket.key);
            if (position == npos || found) {
                throw std::logic_error("flat hash index rebuild invariant violated");
            }
            auto& target = replacement[position];
            target.key.emplace(*bucket.key);
            target.entry = bucket.entry;
            target.state = State::occupied;
        }
        buckets_.swap(replacement);
        tombstones_ = 0;
    }

    void rehash(std::size_t requested_capacity) {
        std::size_t capacity = 16;
        while (capacity < requested_capacity) {
            if (capacity > std::numeric_limits<std::size_t>::max() / 2) {
                throw std::length_error("flat hash index capacity overflow");
            }
            capacity *= 2;
        }

        // Build the complete replacement before touching the live table.  An
        // allocation, hash/equality call, or copy of a copyable Key may throw;
        // in all of those cases the original buckets and counters remain
        // unchanged.  vector::swap is the no-throw commit for std::allocator.
        std::vector<Bucket> replacement(capacity);
        std::size_t replacement_size = 0;
        for (const auto& bucket : buckets_) {
            if (bucket.state != State::occupied) {
                continue;
            }
            const auto [position, found] =
                find_insert_position_in(replacement, *bucket.key);
            if (position == npos) {
                throw std::length_error("flat hash index is full");
            }
            if (found) {
                throw std::logic_error("flat hash index contains duplicate keys");
            }
            auto& target = replacement[position];
            target.key.emplace(*bucket.key);
            target.entry = bucket.entry;
            target.state = State::occupied;
            ++replacement_size;
        }

        buckets_.swap(replacement);
        size_ = replacement_size;
        tombstones_ = 0;
    }

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

} // namespace uc::detail