#pragma once

#include <Core/FxBitset.hpp>
#include <Core/FxHash.hpp>
#include <Core/FxSizedArray.hpp>
#include <map>

/**
 * @brief Cache section that can hold any number of items.
 */
template <typename TItemType>
struct FxItemCacheSection_MultiItem
{
    FxBitset ItemsInUse;
    FxSizedArray<TItemType> Items;

public:
    void Create(uint32 max_items)
    {
        Items.Create(max_items);
        ItemsInUse.InitZero(max_items);
    }

    TItemType* RequestNewItem()
    {
        int index = ItemsInUse.FindNextFreeBit();

        if (index == FxBitset::scNoFreeBits) {
            return nullptr;
        }

        ItemsInUse.Set(index);

        return &Items[index];
    }

    void ReleaseItem(TItemType* item)
    {
        auto item_index = Items.GetItemIndex(item);

        // If the item was found, clear the bit
        if (item_index != FxSizedArray<TItemType>::scItemNotFound) {
            ItemsInUse.Unset(item_index);
        }
    }
};


/**
 * @brief Cache section that holds a single item.
 */
template <typename TItemType>
struct FxItemCacheSection_SingleItem
{
    using ItemType = TItemType;

    TItemType* pItem = nullptr;
    bool bInUse : 1 = 0;

public:
    void Create(uint32 max_items) { pItem = FxMemPool::Alloc<TItemType>(sizeof(TItemType)); }

    TItemType* RequestNewItem()
    {
        if (bInUse) {
            return nullptr;
        }

        bInUse = true;

        return pItem;
    }

    void ReleaseItem(TItemType* item)
    {
        if (pItem != item) {
            return;
        }

        bInUse = false;
    }
};

template <typename T, typename TItemType>
concept C_IsItemCacheSection =
    C_IsAnyBaseOf<T, FxItemCacheSection_MultiItem<TItemType>, FxItemCacheSection_SingleItem<TItemType>>;

template <typename TKeyType, typename TItemType, typename TCacheSectionType>
    requires C_IsItemCacheSection<TCacheSectionType, TItemType>
class FxItemCache
{
public:
    using KeyType = TKeyType;

public:
    FxItemCache() = default;

    void Create(uint32 max_items_per_section) { mItemsPerSection = max_items_per_section; }


    bool IsItemInCache(TKeyType key)
    {
        auto section_it = mCache.find(key);

        return (section_it == mCache.end());
    }

    TItemType* RequestGenericItem(TKeyType key)
    {
        auto section_it = mCache.find(key);

        if (section_it == mCache.end()) {
            CreateCacheSection(key, mItemsPerSection);
            auto section_it = mCache.find(key);

            FxAssert(section_it != mCache.end());
        }

        TCacheSectionType& section = section_it->second;

        return section.RequestNewItem();
    }

    void ReleaseGenericItem(TKeyType key, TItemType* item)
    {
        auto section_it = mCache.find(key);

        if (section_it == mCache.end()) {
            FxLogError("Could not find cache section to free from!");
            return;
        }

        TCacheSectionType& section = section_it->second;

        auto item_index = section.Items.GetItemIndex(item);

        // If the item was found, clear the bit
        if (item_index != FxSizedArray<TItemType>::scItemNotFound) {
            section.ItemsInUse.Unset(item_index);
        }
    }

    void Destroy();

    virtual ~FxItemCache() { Destroy(); }

protected:
    TCacheSectionType& CreateCacheSection(TKeyType key, uint32 number_of_items)
    {
        TCacheSectionType cache_section {};

        cache_section.Items.InitCapacity(number_of_items);
        cache_section.ItemsInUse.InitZero(number_of_items);

        auto iter = mCache.insert({ key, std::move(cache_section) });

        return iter.first->second;
    }

public:
    std::unordered_map<TKeyType, TCacheSectionType> mCache;

protected:
    uint32 mItemsPerSection = 0;
};

/**
 * @brief A temporary handle that is available only for the current scope.
 */
template <typename TKeyType, typename TItemType, typename TCacheType>
struct FxItemCacheHandle
{
    using CacheType = TCacheType;
    using KeyType = TKeyType;

public:
    FxItemCacheHandle(TCacheType* cache, TItemType* item, TKeyType key) : Item(item), Key(key), mCache(cache) {}
    FxItemCacheHandle(const FxItemCacheHandle& other) = delete;
    FxItemCacheHandle(FxItemCacheHandle&& other)
    {
        Item = other.Item;
        Key = other.Key;

        mCache = other.mCache;

        other.Item = nullptr;
    }

    /**
     * @brief Releases the handle before the destructor is called.
     */
    FX_FORCE_INLINE void Release()
    {
        if (Item == nullptr) {
            return;
        }

        mCache->ReleaseGenericItem(Key, Item);
    }

    ~FxItemCacheHandle() { Release(); }

    TItemType* operator->() { return Item; }
    TItemType& operator*() { return *Item; }

public:
    TItemType* Item = nullptr;
    TKeyType Key;

private:
    CacheType* mCache;
};