/*********************************************************************
 * \file   AssetContainer.hpp
 *
 * \author Romain BOULLARD
 * \date   May 2026
 *********************************************************************/
#ifndef BIGFOOT_ENGINE_ASSET_ASSETCONTAINER_HPP
#define BIGFOOT_ENGINE_ASSET_ASSETCONTAINER_HPP
#include <Engine/Asset/Asset.hpp>
#include <Engine/EngineAssertHandler.hpp>

#include <Utils/Containers/SlotMap.hpp>

#include <ankerl/unordered_dense.h>

#include <cstdint>

namespace Bigfoot
{
class ContainerHandle
{
  public:
    ContainerHandle() = default;

    ContainerHandle(const ContainerHandle& p_handle) = default;
    ContainerHandle(ContainerHandle&& p_handle) = default;

    ~ContainerHandle() = default;

    /*
     * Increment the HardRef count
     *
     */
    void IncrementHardRefCount()
    {
        ++m_hardRefCount;
    }

    /*
     * Decrement the HardRef count
     *
     */
    void DecrementHardRefCount()
    {
        --m_hardRefCount;
    }

    /*
     * Increment the SoftRef count
     *
     */
    void IncrementSoftRefCount()
    {
        ++m_softRefCount;
    }

    /*
     * Decrement the SoftRef count
     *
     */
    void DecrementSoftRefCount()
    {
        --m_softRefCount;
    }

    /*
     * Get the Asset
     *
     * \return The asset, can be nullptr
     */
    [[nodiscard]]
    AssetBase* GetAsset()
    {
        return m_asset;
    }

    /*
     * Get the Asset
     *
     * \return The asset
     */
    [[nodiscard]]
    const AssetBase* GetAsset() const
    {
        return m_asset;
    }

    /*
     * Set the asset
     *
     * \param p_asset The asset
     */
    void SetAsset(AssetBase* p_asset)
    {
        m_asset = p_asset;
    }

    ContainerHandle& operator=(const ContainerHandle& p_handle) = default;
    ContainerHandle& operator=(ContainerHandle&& p_handle) = default;

  private:
    /*
     * The HardRef count
     */
    std::uint32_t m_hardRefCount = 0;

    /*
     * The SoftRef count
     */
    std::uint32_t m_softRefCount = 0;

    /*
     * The asset
     */
    AssetBase* m_asset = nullptr;
};

/****************************************************************************************/

class AssetContainerBase
{
  public:
    virtual ~AssetContainerBase() = default;

    /*
     * Check if an Asset is in the AssetContainer or not
     *
     * \param p_uuid The UUID to check
     * \return True if in AssetContainer, false otherwise
     */
    [[nodiscard]]
    bool Has(const UUID& p_uuid) const
    {
        const Entry* entry = FindEntry(p_uuid);
        return entry != nullptr && entry->m_assetSlotKey.Valid();
    }

    /*
     * Acquires a ContainerHandle for a given UUID
     * Will create it if it does not exists
     *
     * \param p_uuid The UUID to get ContainerHandle for
     * \return A ContainerHandle for the UUID
     */
    [[nodiscard]]
    ContainerHandle& AcquireContainerHandle(const UUID& p_uuid)
    {
        return *m_handles.Get(GetOrCreateEntry(p_uuid).m_handleSlotKey);
    }

    /*
     * Remove an Asset from the AssetContainer
     *
     * \param p_uuid The UUID to remove
     */
    virtual void Remove(const UUID& p_uuid)
    {
        if (Entry* entry = FindEntry(p_uuid); entry != nullptr && entry->m_handleSlotKey.Valid())
        {
            m_handles.Remove(entry->m_handleSlotKey);
            m_uuidToEntry.erase(p_uuid);
        }
    }

    AssetContainerBase& operator=(const AssetContainerBase& p_container) = delete;
    AssetContainerBase& operator=(AssetContainerBase&& p_container) = default;

  protected:
    struct Entry
    {
        /*
         * SlotMapKey to the handle
         */
        SlotMapKey m_handleSlotKey;

        /*
         * SlotMapKey to the asset
         */
        SlotMapKey m_assetSlotKey;
    };

    AssetContainerBase() = default;

    AssetContainerBase(const AssetContainerBase& p_container) = delete;
    AssetContainerBase(AssetContainerBase&& p_container) = default;

    /*
     * Finds an entry matching the Asset
     *
     * \param p_uuid The UUID to find
     * \return Pointer to an Entry, nullptr otherwise
     */
    [[nodiscard]]
    Entry* FindEntry(const UUID& p_uuid)
    {
        const auto it = m_uuidToEntry.find(p_uuid);
        return it != m_uuidToEntry.end() ? &it->second : nullptr;
    }

    /*
     * Finds an entry matching the Asset
     *
     * \param p_uuid The UUID to find
     * \return Pointer to an Entry, nullptr otherwise
     */
    [[nodiscard]]
    const Entry* FindEntry(const UUID& p_uuid) const
    {
        const auto it = m_uuidToEntry.find(p_uuid);
        return it != m_uuidToEntry.end() ? &it->second : nullptr;
    }

    /*
     * Finds an entry matching the Asset
     * Creates it if it does not exist
     *
     * \param p_uuid The UUID to find
     * \return The entry matching the Asset
     */
    [[nodiscard]]
    Entry& GetOrCreateEntry(const UUID& p_uuid)
    {
        auto [it, inserted] = m_uuidToEntry.try_emplace(p_uuid);
        if (inserted)
        {
            it->second.m_handleSlotKey = m_handles.Insert();
        }

        return it->second;
    }

    /*
     * Get a ContainerHandle
     *
     * \param p_handleSlotKey The key to the container
     * \return The container
     */
    [[nodiscard]]
    ContainerHandle& FindContainerHandle(const SlotMapKey p_handleSlotKey)
    {
        return *m_handles.Get(p_handleSlotKey);
    }

  private:
    /*
     * HandleContainers
     */
    StableSlotMap<ContainerHandle> m_handles;

    /*
     * Every Asset added to the map
     */
    ankerl::unordered_dense::segmented_map<UUID, Entry> m_uuidToEntry;
};

/****************************************************************************************/

template<class ASSET>
class AssetContainer: public AssetContainerBase
{
  public:
    AssetContainer() = default;

    AssetContainer(const AssetContainer& p_container) = delete;
    AssetContainer(AssetContainer&& p_container) = default;

    ~AssetContainer() = default;

    /*
     * Inserts an Asset to the AssetContainer
     *
     * \tparam ARGS Arguments to construct the asset
     * \param p_args The arguments
     * \return The key to the inserted asset.
     */
    template<class... ARGS>
    [[nodiscard]]
    SlotMapKey Insert(ARGS&&... p_args)
    {
        SlotMapKey slotKey = m_assets.Insert(std::forward<ARGS>(p_args)...);
        ASSET* asset = m_assets.Get(slotKey);

        Entry& entry = GetOrCreateEntry(asset->GetHeader().uuid);
        if (entry.m_assetSlotKey.Valid())
        {
            m_assets.Remove(slotKey);
            slotKey = entry.m_assetSlotKey;
        }
        else
        {
            entry.m_assetSlotKey = slotKey;
            FindContainerHandle(entry.m_handleSlotKey).SetAsset(asset);
        }

        return slotKey;
    }

    /*
     * Get the Asset
     *
     * \param p_key The key to the asset
     * \return Pointer to the Asset
     */
    [[nodiscard]]
    ASSET* Get(const SlotMapKey p_key)
    {
        return m_assets.Get(p_key);
    }

    /*
     * Get the Asset
     *
     * \param p_key The key to the asset
     * \return Pointer to the Asset
     */
    [[nodiscard]]
    const ASSET* Get(const SlotMapKey p_key) const
    {
        return m_assets.Get(p_key);
    }

    /*
     * Remove an Asset from the AssetContainer
     *
     * \param p_uuid The UUID to remove
     */
    void Remove(const UUID& p_uuid) override
    {
        if (Entry* entry = FindEntry(p_uuid); entry != nullptr && entry->m_assetSlotKey.Valid())
        {
            FindContainerHandle(entry->m_handleSlotKey).SetAsset(nullptr);
            m_assets.Remove(entry->m_assetSlotKey);
            entry->m_assetSlotKey = SlotMapKey {};
        }

        AssetContainerBase::Remove(p_uuid);
    }

    AssetContainer& operator=(const AssetContainer& p_container) = delete;
    AssetContainer& operator=(AssetContainer&& p_container) = default;

  private:
    /*
     * The assets
     */
    StableSlotMap<ASSET> m_assets;
};
}   // namespace Bigfoot

#endif
