/*********************************************************************
 * \file   SlotMap.hpp
 *
 * \author Romain BOULLARD
 * \date   May 2026
 *********************************************************************/
#ifndef BIGFOOT_UTILS_CONTAINERS_SLOTMAP_HPP
#define BIGFOOT_UTILS_CONTAINERS_SLOTMAP_HPP
#include <Utils/UtilsAssertHandler.hpp>

#include <EASTL/array.h>
#include <EASTL/bitset.h>
#include <EASTL/iterator.h>
#include <EASTL/unique_ptr.h>
#include <EASTL/vector.h>

#include <cstddef>
#include <cstdint>
#include <limits>
#include <memory>
#include <new>

namespace Bigfoot
{
class SlotMapKey
{
  public:
    using IndexType = std::uint32_t;
    using VersionType = std::uint32_t;

  private:
    static constexpr std::uint32_t VERSION_BIT_COUNT = sizeof(VersionType) * std::numeric_limits<unsigned char>::digits;
    static constexpr std::uint32_t INDEX_BIT_COUNT = sizeof(IndexType) * std::numeric_limits<unsigned char>::digits;
    static_assert(VERSION_BIT_COUNT + INDEX_BIT_COUNT <= 64,
                  "We cant construct a 64 bit key from the given Version and Index types!");

    static constexpr std::uint64_t INDEX_MASK = (static_cast<std::uint64_t>(1) << INDEX_BIT_COUNT) - 1;

  public:
    static constexpr VersionType INVALID_VERSION = 0;

    static constexpr IndexType MAX_INDEX = std::numeric_limits<IndexType>::max();

    constexpr SlotMapKey(const VersionType p_version, const IndexType p_index):
        m_key((static_cast<std::uint64_t>(p_version) << INDEX_BIT_COUNT) |
              (static_cast<std::uint64_t>(p_index) & INDEX_MASK))
    {
    }

    constexpr SlotMapKey():
        SlotMapKey(INVALID_VERSION, 0)
    {
    }

    constexpr SlotMapKey(const SlotMapKey& p_slotKey) = default;
    constexpr SlotMapKey(SlotMapKey&& p_slotKey) = default;

    constexpr ~SlotMapKey() = default;

    constexpr bool Valid() const
    {
        return GetVersion() != INVALID_VERSION;
    }

    constexpr VersionType GetVersion() const
    {
        return static_cast<VersionType>(m_key >> INDEX_BIT_COUNT);
    }

    constexpr IndexType GetIndex() const
    {
        return static_cast<IndexType>(m_key & INDEX_MASK);
    }

    constexpr SlotMapKey& operator=(const SlotMapKey& p_slotKey) = default;
    constexpr SlotMapKey& operator=(SlotMapKey&& p_slotKey) = default;

    [[nodiscard]]
    constexpr bool operator==(const SlotMapKey& p_other) const = default;

  private:
    std::uint64_t m_key;
};

class SlotMapBase
{
  public:
    ~SlotMapBase() = default;

    [[nodiscard]]
    bool Has(const SlotMapKey p_slotKey) const
    {
        return p_slotKey.Valid() && (p_slotKey.GetIndex() < m_slots.size() &&
                                     p_slotKey.GetVersion() == m_slots[p_slotKey.GetIndex()].GetVersion());
    }

    [[nodiscard]]
    std::size_t Size() const
    {
        return m_size;
    }

    [[nodiscard]]
    bool Empty() const
    {
        return m_size == 0;
    }

    SlotMapBase& operator=(const SlotMapBase& p_base) = default;
    SlotMapBase& operator=(SlotMapBase&& p_base) = default;

  protected:
    SlotMapBase() = default;

    SlotMapBase(const SlotMapBase& p_base) = default;
    SlotMapBase(SlotMapBase&& p_base) = default;

    [[nodiscard]]
    bool HasFreeSlots() const
    {
        return m_freeSlotHead != std::numeric_limits<SlotMapKey::IndexType>::max();
    }

    void RecycleSlot(const SlotMapKey::VersionType p_version, const SlotMapKey::IndexType p_slotIndex)
    {
        if (p_version == SlotMapKey::INVALID_VERSION)
        {
            m_slots[p_slotIndex] = SlotMapKey {SlotMapKey::INVALID_VERSION, 0};
            return;
        }

        m_slots[p_slotIndex] = SlotMapKey {p_version, m_freeSlotHead};
        m_freeSlotHead = p_slotIndex;
    }

    void ResetSlots()
    {
        m_size = 0;
        m_slots.clear();
        m_freeSlotHead = std::numeric_limits<SlotMapKey::IndexType>::max();
    }

    void ReserveSlots(const std::size_t p_size)
    {
        m_slots.reserve(p_size);
    }

    std::size_t m_size = 0;
    eastl::vector<SlotMapKey> m_slots;
    SlotMapKey::IndexType m_freeSlotHead = std::numeric_limits<SlotMapKey::IndexType>::max();
};

template<class TYPE>
class DenseSlotMap: public SlotMapBase
{
  public:
    using iterator = typename eastl::vector<TYPE>::iterator;
    using const_iterator = typename eastl::vector<TYPE>::const_iterator;
    using reverse_iterator = typename eastl::vector<TYPE>::reverse_iterator;
    using const_reverse_iterator = typename eastl::vector<TYPE>::const_reverse_iterator;

    DenseSlotMap() = default;

    DenseSlotMap(const DenseSlotMap& p_map) = default;
    DenseSlotMap(DenseSlotMap&& p_map) = default;

    ~DenseSlotMap() = default;

    template<class... ARGS>
    [[nodiscard]]
    SlotMapKey Insert(ARGS&&... p_args)
    {
        const SlotMapKey::IndexType dataIndex = static_cast<SlotMapKey::IndexType>(m_size);
        m_data.emplace_back(std::forward<ARGS>(p_args)...);
        ++m_size;

        return AcquireSlot(dataIndex);
    }

    void Remove(const SlotMapKey p_slotKey)
    {
        if (!Has(p_slotKey))
        {
            return;
        }

        const SlotMapKey::IndexType dataIndex = GetDataIndex(p_slotKey);

        if (dataIndex != m_size - 1)
        {
            m_data[dataIndex] = eastl::move(m_data.back());
        }
        m_data.pop_back();
        --m_size;

        ReleaseSlot(p_slotKey, dataIndex);
    }

    [[nodiscard]]
    TYPE* Get(const SlotMapKey p_slotKey)
    {
        return Has(p_slotKey) ? &m_data[GetDataIndex(p_slotKey)] : nullptr;
    }

    [[nodiscard]]
    const TYPE* Get(const SlotMapKey p_slotKey) const
    {
        return Has(p_slotKey) ? &m_data[GetDataIndex(p_slotKey)] : nullptr;
    }

    void Reset()
    {
        m_data.clear();
        m_dataToSlots.clear();
        ResetSlots();
    }

    void Clear()
    {
        for (const std::size_t slotIndex: m_dataToSlots)
        {
            RecycleSlot(m_slots[slotIndex].GetVersion() + 1, static_cast<SlotMapKey::IndexType>(slotIndex));
        }

        m_size = 0;
        m_data.clear();
        m_dataToSlots.clear();
    }

    [[nodiscard]]
    std::size_t Capacity() const
    {
        return m_data.capacity();
    }

    void Reserve(const std::size_t p_size)
    {
        m_data.reserve(p_size);
        m_dataToSlots.reserve(p_size);
        ReserveSlots(p_size);
    }

    [[nodiscard]]
    iterator begin()
    {
        return m_data.begin();
    }

    [[nodiscard]]
    iterator end()
    {
        return m_data.end();
    }

    [[nodiscard]]
    const_iterator begin() const
    {
        return m_data.begin();
    }

    [[nodiscard]]
    const_iterator end() const
    {
        return m_data.end();
    }

    [[nodiscard]]
    reverse_iterator rbegin()
    {
        return m_data.rbegin();
    }

    [[nodiscard]]
    reverse_iterator rend()
    {
        return m_data.rend();
    }

    [[nodiscard]]
    const_reverse_iterator rbegin() const
    {
        return m_data.rbegin();
    }

    [[nodiscard]]
    const_reverse_iterator rend() const
    {
        return m_data.rend();
    }

    [[nodiscard]]
    const_iterator cbegin() const
    {
        return m_data.cbegin();
    }

    [[nodiscard]]
    const_iterator cend() const
    {
        return m_data.cend();
    }

    [[nodiscard]]
    const_reverse_iterator crbegin() const
    {
        return m_data.crbegin();
    }

    [[nodiscard]]
    const_reverse_iterator crend() const
    {
        return m_data.crend();
    }

    DenseSlotMap& operator=(const DenseSlotMap& p_map) = default;
    DenseSlotMap& operator=(DenseSlotMap&& p_map) = default;

  private:
    [[nodiscard]]
    SlotMapKey::IndexType GetDataIndex(const SlotMapKey p_slotKey) const
    {
        return m_slots[p_slotKey.GetIndex()].GetIndex();
    }

    [[nodiscard]]
    SlotMapKey AcquireSlot(const SlotMapKey::IndexType p_dataIndex)
    {
        ASSERT(UtilsAssertHandler, m_slots.size() < SlotMapKey::MAX_INDEX, "All slots have been exhausted!");

        if (HasFreeSlots())
        {
            const SlotMapKey::IndexType slotIndex = m_freeSlotHead;

            m_freeSlotHead = m_slots[slotIndex].GetIndex();

            m_slots[slotIndex] = SlotMapKey {m_slots[slotIndex].GetVersion(), p_dataIndex};
            m_dataToSlots.push_back(slotIndex);

            return SlotMapKey {m_slots[slotIndex].GetVersion(), slotIndex};
        }

        const SlotMapKey::IndexType slotIndex = static_cast<SlotMapKey::IndexType>(m_slots.size());
        m_slots.emplace_back(1, p_dataIndex);
        m_dataToSlots.push_back(slotIndex);

        return SlotMapKey {1, slotIndex};
    }

    void ReleaseSlot(const SlotMapKey p_slotKey, const SlotMapKey::IndexType p_dataIndex)
    {
        m_dataToSlots.erase_unsorted(m_dataToSlots.begin() + p_dataIndex);

        if (p_dataIndex < m_size)
        {
            m_slots[m_dataToSlots[p_dataIndex]] = SlotMapKey {m_slots[m_dataToSlots[p_dataIndex]].GetVersion(),
                                                              p_dataIndex};
        }

        RecycleSlot(p_slotKey.GetVersion() + 1, p_slotKey.GetIndex());
    }

    eastl::vector<TYPE> m_data;
    eastl::vector<std::size_t> m_dataToSlots;
};

template<class TYPE>
class StableSlotMap: public SlotMapBase
{
  private:
    static constexpr std::size_t PAGE_SIZE = 64;

    struct AlignedSlot
    {
        alignas(TYPE) std::byte m_bytes[sizeof(TYPE)];
    };

    struct Page
    {
        Page() = default;

        Page(const Page& p_page) = delete;
        Page(Page&& p_page) = delete;

        ~Page()
        {
            for (std::size_t offset = m_occupied.find_first(); offset != PAGE_SIZE; offset = m_occupied.find_next(offset))
            {
                std::destroy_at(SlotAt(offset));
            }
        }

        Page& operator=(const Page& p_page) = delete;
        Page& operator=(Page&& p_page) = delete;

        [[nodiscard]]
        TYPE* SlotAt(const std::size_t p_offset)
        {
            return std::launder(reinterpret_cast<TYPE*>(m_storage[p_offset].m_bytes));
        }

        [[nodiscard]]
        const TYPE* SlotAt(const std::size_t p_offset) const
        {
            return std::launder(reinterpret_cast<const TYPE*>(m_storage[p_offset].m_bytes));
        }

        eastl::array<AlignedSlot, PAGE_SIZE> m_storage;
        eastl::bitset<PAGE_SIZE> m_occupied;
    };

  public:
    class iterator
    {
      public:
        using iterator_category = eastl::bidirectional_iterator_tag;
        using value_type = TYPE;
        using difference_type = std::ptrdiff_t;
        using pointer = TYPE*;
        using reference = TYPE&;

        iterator() = default;

        [[nodiscard]]
        reference operator*() const
        {
            return *m_current;
        }

        [[nodiscard]]
        pointer operator->() const
        {
            return m_current;
        }

        iterator& operator++()
        {
            Advance();
            return *this;
        }

        iterator operator++(int)
        {
            const iterator temp = *this;
            ++(*this);
            return temp;
        }

        iterator& operator--()
        {
            Retreat();
            return *this;
        }

        iterator operator--(int)
        {
            const iterator temp = *this;
            --(*this);
            return temp;
        }

        [[nodiscard]]
        bool operator==(const iterator& p_other) const
        {
            return m_map == p_other.m_map && m_current == p_other.m_current;
        }

      private:
        friend class StableSlotMap;

        explicit iterator(StableSlotMap* p_map):
            m_map(p_map)
        {
        }

        void SeekFirstFrom(std::size_t p_pageIndex)
        {
            for (; p_pageIndex < m_map->m_pages.size(); ++p_pageIndex)
            {
                Page* const page = m_map->m_pages[p_pageIndex].get();
                const std::size_t offset = page->m_occupied.find_first();
                if (offset != PAGE_SIZE)
                {
                    m_pageIndex = p_pageIndex;
                    m_page = page;
                    m_offset = offset;
                    m_current = page->SlotAt(offset);
                    return;
                }
            }

            m_pageIndex = m_map->m_pages.size();
            m_page = nullptr;
            m_offset = 0;
            m_current = nullptr;
        }

        void SeekLastFrom(std::size_t p_pageIndex)
        {
            for (;;)
            {
                Page* const page = m_map->m_pages[p_pageIndex].get();
                const std::size_t offset = page->m_occupied.find_last();
                if (offset != PAGE_SIZE)
                {
                    m_pageIndex = p_pageIndex;
                    m_page = page;
                    m_offset = offset;
                    m_current = page->SlotAt(offset);
                    return;
                }

                if (p_pageIndex == 0)
                {
                    return;
                }
                --p_pageIndex;
            }
        }

        void Advance()
        {
            const std::size_t offset = m_page->m_occupied.find_next(m_offset);
            if (offset != PAGE_SIZE)
            {
                m_offset = offset;
                m_current = m_page->SlotAt(offset);
                return;
            }

            SeekFirstFrom(m_pageIndex + 1);
        }

        void Retreat()
        {
            if (m_current == nullptr)
            {
                SeekLastFrom(m_map->m_pages.size() - 1);
                return;
            }

            const std::size_t offset = m_page->m_occupied.find_prev(m_offset);
            if (offset != PAGE_SIZE)
            {
                m_offset = offset;
                m_current = m_page->SlotAt(offset);
                return;
            }

            if (m_pageIndex == 0)
            {
                return;
            }
            SeekLastFrom(m_pageIndex - 1);
        }

        StableSlotMap* m_map = nullptr;
        Page* m_page = nullptr;
        std::size_t m_pageIndex = 0;
        std::size_t m_offset = 0;
        TYPE* m_current = nullptr;
    };

    class const_iterator
    {
      public:
        using iterator_category = eastl::bidirectional_iterator_tag;
        using value_type = TYPE;
        using difference_type = std::ptrdiff_t;
        using pointer = const TYPE*;
        using reference = const TYPE&;

        const_iterator() = default;

        [[nodiscard]]
        reference operator*() const
        {
            return *m_current;
        }

        [[nodiscard]]
        pointer operator->() const
        {
            return m_current;
        }

        const_iterator& operator++()
        {
            Advance();
            return *this;
        }

        const_iterator operator++(int)
        {
            const const_iterator temp = *this;
            ++(*this);
            return temp;
        }

        const_iterator& operator--()
        {
            Retreat();
            return *this;
        }

        const_iterator operator--(int)
        {
            const const_iterator temp = *this;
            --(*this);
            return temp;
        }

        [[nodiscard]]
        bool operator==(const const_iterator& p_other) const
        {
            return m_map == p_other.m_map && m_current == p_other.m_current;
        }

      private:
        friend class StableSlotMap;

        explicit const_iterator(const StableSlotMap* p_map):
            m_map(p_map)
        {
        }

        void SeekFirstFrom(std::size_t p_pageIndex)
        {
            for (; p_pageIndex < m_map->m_pages.size(); ++p_pageIndex)
            {
                const Page* const page = m_map->m_pages[p_pageIndex].get();
                const std::size_t offset = page->m_occupied.find_first();
                if (offset != PAGE_SIZE)
                {
                    m_pageIndex = p_pageIndex;
                    m_page = page;
                    m_offset = offset;
                    m_current = page->SlotAt(offset);
                    return;
                }
            }

            m_pageIndex = m_map->m_pages.size();
            m_page = nullptr;
            m_offset = 0;
            m_current = nullptr;
        }

        void SeekLastFrom(std::size_t p_pageIndex)
        {
            for (;;)
            {
                const Page* const page = m_map->m_pages[p_pageIndex].get();
                const std::size_t offset = page->m_occupied.find_last();
                if (offset != PAGE_SIZE)
                {
                    m_pageIndex = p_pageIndex;
                    m_page = page;
                    m_offset = offset;
                    m_current = page->SlotAt(offset);
                    return;
                }

                if (p_pageIndex == 0)
                {
                    return;
                }
                --p_pageIndex;
            }
        }

        void Advance()
        {
            const std::size_t offset = m_page->m_occupied.find_next(m_offset);
            if (offset != PAGE_SIZE)
            {
                m_offset = offset;
                m_current = m_page->SlotAt(offset);
                return;
            }

            SeekFirstFrom(m_pageIndex + 1);
        }

        void Retreat()
        {
            if (m_current == nullptr)
            {
                SeekLastFrom(m_map->m_pages.size() - 1);
                return;
            }

            const std::size_t offset = m_page->m_occupied.find_prev(m_offset);
            if (offset != PAGE_SIZE)
            {
                m_offset = offset;
                m_current = m_page->SlotAt(offset);
                return;
            }

            if (m_pageIndex == 0)
            {
                return;
            }
            SeekLastFrom(m_pageIndex - 1);
        }

        const StableSlotMap* m_map = nullptr;
        const Page* m_page = nullptr;
        std::size_t m_pageIndex = 0;
        std::size_t m_offset = 0;
        const TYPE* m_current = nullptr;
    };

    using reverse_iterator = eastl::reverse_iterator<iterator>;
    using const_reverse_iterator = eastl::reverse_iterator<const_iterator>;

    StableSlotMap() = default;

    StableSlotMap(const StableSlotMap& p_map) = delete;
    StableSlotMap(StableSlotMap&& p_map) = default;

    ~StableSlotMap() = default;

    template<class... ARGS>
    [[nodiscard]]
    SlotMapKey Insert(ARGS&&... p_args)
    {
        const SlotMapKey key = AcquireFreeSlot();
        const SlotMapKey::IndexType index = key.GetIndex();
        const std::size_t pageIndex = index / PAGE_SIZE;
        if (pageIndex >= m_pages.size())
        {
            m_pages.push_back(eastl::make_unique<Page>());
        }

        Page& page = *m_pages[pageIndex];
        const std::size_t offset = index % PAGE_SIZE;
        ::new (static_cast<void*>(page.m_storage[offset].m_bytes)) TYPE(std::forward<ARGS>(p_args)...);
        page.m_occupied.set(offset);
        ++m_size;

        return key;
    }

    void Remove(const SlotMapKey p_slotKey)
    {
        if (!Has(p_slotKey))
        {
            return;
        }

        const SlotMapKey::IndexType index = p_slotKey.GetIndex();
        Page& page = *m_pages[index / PAGE_SIZE];
        const std::size_t offset = index % PAGE_SIZE;
        std::destroy_at(page.SlotAt(offset));
        page.m_occupied.reset(offset);
        --m_size;

        ReleaseSlotVersion(index);
    }

    [[nodiscard]]
    TYPE* Get(const SlotMapKey p_slotKey)
    {
        return Has(p_slotKey) ? &DataAt(p_slotKey.GetIndex()) : nullptr;
    }

    [[nodiscard]]
    const TYPE* Get(const SlotMapKey p_slotKey) const
    {
        return Has(p_slotKey) ? &DataAt(p_slotKey.GetIndex()) : nullptr;
    }

    void Reset()
    {
        for (eastl::unique_ptr<Page>& page: m_pages)
        {
            page = eastl::make_unique<Page>();
        }

        ResetSlots();
    }

    void Clear()
    {
        for (std::size_t pageIndex = 0; pageIndex < m_pages.size(); ++pageIndex)
        {
            Page& page = *m_pages[pageIndex];
            for (std::size_t offset = page.m_occupied.find_first(); offset != PAGE_SIZE;
                 offset = page.m_occupied.find_next(offset))
            {
                std::destroy_at(page.SlotAt(offset));
                ReleaseSlotVersion(static_cast<SlotMapKey::IndexType>(pageIndex * PAGE_SIZE + offset));
            }
            page.m_occupied.reset();
        }

        m_size = 0;
    }

    [[nodiscard]]
    std::size_t Capacity() const
    {
        return m_pages.size() * PAGE_SIZE;
    }

    void Reserve(const std::size_t p_size)
    {
        const std::size_t pagesNeeded = (p_size + PAGE_SIZE - 1) / PAGE_SIZE;
        m_pages.reserve(pagesNeeded);
        for (std::size_t pageIndex = m_pages.size(); pageIndex < pagesNeeded; ++pageIndex)
        {
            m_pages.push_back(eastl::make_unique<Page>());
        }

        ReserveSlots(p_size);
    }

    [[nodiscard]]
    iterator begin()
    {
        iterator it {this};
        it.SeekFirstFrom(0);
        return it;
    }

    [[nodiscard]]
    iterator end()
    {
        return iterator {this};
    }

    [[nodiscard]]
    const_iterator begin() const
    {
        const_iterator it {this};
        it.SeekFirstFrom(0);
        return it;
    }

    [[nodiscard]]
    const_iterator end() const
    {
        return const_iterator {this};
    }

    [[nodiscard]]
    reverse_iterator rbegin()
    {
        return reverse_iterator {end()};
    }

    [[nodiscard]]
    reverse_iterator rend()
    {
        return reverse_iterator {begin()};
    }

    [[nodiscard]]
    const_reverse_iterator rbegin() const
    {
        return const_reverse_iterator {end()};
    }

    [[nodiscard]]
    const_reverse_iterator rend() const
    {
        return const_reverse_iterator {begin()};
    }

    [[nodiscard]]
    const_iterator cbegin() const
    {
        return begin();
    }

    [[nodiscard]]
    const_iterator cend() const
    {
        return end();
    }

    [[nodiscard]]
    const_reverse_iterator crbegin() const
    {
        return rbegin();
    }

    [[nodiscard]]
    const_reverse_iterator crend() const
    {
        return rend();
    }

    StableSlotMap& operator=(const StableSlotMap& p_map) = delete;
    StableSlotMap& operator=(StableSlotMap&& p_map) = default;

  private:
    [[nodiscard]]
    TYPE& DataAt(const std::size_t p_index)
    {
        return *m_pages[p_index / PAGE_SIZE]->SlotAt(p_index % PAGE_SIZE);
    }

    [[nodiscard]]
    const TYPE& DataAt(const std::size_t p_index) const
    {
        return *m_pages[p_index / PAGE_SIZE]->SlotAt(p_index % PAGE_SIZE);
    }

    [[nodiscard]]
    SlotMapKey AcquireFreeSlot()
    {
        ASSERT(UtilsAssertHandler, m_slots.size() < SlotMapKey::MAX_INDEX, "All slots have been exhausted!");

        if (HasFreeSlots())
        {
            const SlotMapKey::IndexType slotIndex = m_freeSlotHead;
            m_freeSlotHead = m_slots[slotIndex].GetIndex();
            m_slots[slotIndex] = SlotMapKey {m_slots[slotIndex].GetVersion(), slotIndex};
            return m_slots[slotIndex];
        }

        const SlotMapKey::IndexType slotIndex = static_cast<SlotMapKey::IndexType>(m_slots.size());
        m_slots.emplace_back(1, slotIndex);
        return m_slots[slotIndex];
    }

    void ReleaseSlotVersion(const SlotMapKey::IndexType p_index)
    {
        RecycleSlot(m_slots[p_index].GetVersion() + 1, p_index);
    }

    eastl::vector<eastl::unique_ptr<Page>> m_pages;
};
}   // namespace Bigfoot

#endif