#pragma once
#include "Upscaler.h"

#include <vector>
#include <queue>
#include <mutex>

template<typename TFeature> class UpscalerGraphicsDevice: public Upscaler
{
public:
    void Shutdown() override
    {
        // Called by AMDUP_ShutdownAPI, destroys all features, cleans up internal resources
        size_t numFeatures = 0;
        {
            std::lock_guard<std::mutex> lock(m_Mutex);
            numFeatures = m_Features.size();
        }

        for (uint32_t slot = 0; slot < numFeatures; ++slot)
        {
            DestroyFeature(slot);
        }

        std::lock_guard<std::mutex> lock(m_Mutex);

        DoShutdown();
    }

    uint32_t CreateFeatureSlot() override
    {
        std::lock_guard<std::mutex> lock(m_Mutex);
        return AllocateFeatureSlot();
    }

    bool InitFeature(const FSR3CommandInitializationData* initData) override
    {
        std::lock_guard<std::mutex> lock(m_Mutex);

        // Called by FSR3PluginEvent::eInit event
        if (initData->featureSlot < 0 || initData->featureSlot >= m_Features.size())
            return false;

        auto& feature = m_Features[initData->featureSlot];
        return InitFeature(feature, initData);
    }

    void DestroyFeature(uint32_t featureSlot) override
    {
        // Called by FSR3PluginEvent::eDestroyFeature event and Shutdown
        uint64_t dispatchFrameValue = 0;
        {
            // We need to lock the features list while we're accessing it, but we also can't hold the lock while blocking the main/render thread at the same time.
            // Otherwise we're very likely to create a deadlock. So instead we only lock for as long as is necessary to grab the data we need from the feature.
            std::lock_guard<std::mutex> lock(m_Mutex);

            if (featureSlot < 0 || featureSlot >= m_Features.size())
                return;

            auto& feature = m_Features[featureSlot];
            if (!IsValidFeature(feature))
                return;

            dispatchFrameValue = feature.dispatchFrameValue;
        }

        // If there's still an upscale dispatch executing on the current frame, wait until rendering is finished before destroying its context.
        AwaitEndOfFrame(dispatchFrameValue);

        std::lock_guard<std::mutex> lock(m_Mutex);

        DestroyContext(m_Features[featureSlot]);
        FreeFeatureSlot(featureSlot);
    }

    void ClearTextureTable(uint32_t featureSlot) override
    {
        if (featureSlot < 0 || featureSlot >= m_Features.size())
            return;

        memset(&m_Features[featureSlot].textureTable, 0, sizeof(FSR3TextureTable));
    }

    void SetTextureSlot(uint32_t featureSlot, uint32_t textureSlot, UnityTextureID textureID, uint32_t width, uint32_t height) override
    {
        if (featureSlot < 0 || featureSlot >= m_Features.size())
            return;

        // We organized the texture table struct to be ordered the same as the texture slot enum
        // This way we can use the texture slot value simply as a pointer offset into the texture table
        FSR3TextureDesc* textureDesc = ((FSR3TextureDesc*)&m_Features[featureSlot].textureTable) + textureSlot;
        SetTexture((FSR3Texture)textureSlot, textureID, width, height, textureDesc);
    }

    void Execute(const FSR3CommandExecutionData* execData) override
    {
        std::lock_guard<std::mutex> lock(m_Mutex);

        // Called by FSR3PluginEvent::eExecute event
        if (execData->featureSlot < 0 || execData->featureSlot >= m_Features.size())
            return;

        auto& feature = m_Features[execData->featureSlot];
        if (!IsValidFeature(feature))
            return;

        Execute(feature, execData);
    }

    void PostExecute(const FSR3CommandExecutionData* execData) override
    {
        std::lock_guard<std::mutex> lock(m_Mutex);

        // Called by FSR3PluginEvent::ePostExecute event
        if (execData->featureSlot < 0 || execData->featureSlot >= m_Features.size())
            return;

        auto& feature = m_Features[execData->featureSlot];
        if (!IsValidFeature(feature))
            return;

        PostExecute(feature, execData);
    }

protected:
    virtual bool IsValidFeature(TFeature& feature) = 0;

    virtual bool InitFeature(TFeature& feature, const FSR3CommandInitializationData* initData) = 0;
    virtual void SetTexture(FSR3Texture textureType, UnityTextureID textureID, uint32_t width, uint32_t height, FSR3TextureDesc* outTextureDesc) = 0;
    virtual void Execute(TFeature& feature, const FSR3CommandExecutionData* execData) = 0;
    virtual void PostExecute(TFeature& feature, const FSR3CommandExecutionData* execData) { }

    virtual void AwaitEndOfFrame(uint64_t frameValue) = 0;
    virtual void DestroyContext(TFeature& feature) = 0;

    virtual void DoShutdown() = 0;

private:
    uint32_t AllocateFeatureSlot()
    {
        if (m_FeatureSlots.empty())
        {
            // Create a new feature if there are no free slots
            uint32_t featureSlot = (uint32_t)m_Features.size();
            m_Features.push_back(std::move(TFeature()));
            memset(&m_Features[featureSlot], 0, sizeof(TFeature));
            return featureSlot;
        }

        // Reallocate an existing free slot
        uint32_t featureSlot = m_FeatureSlots.front();
        m_FeatureSlots.pop();
        return featureSlot;
    }

    void FreeFeatureSlot(uint32_t featureSlot)
    {
        // Prevent duplicate free slots
        auto& slots = m_FeatureSlots._Get_container();
        if (std::find(slots.begin(), slots.end(), featureSlot) != slots.end())
            return;

        m_FeatureSlots.push(featureSlot);
        memset(&m_Features[featureSlot], 0, sizeof(TFeature));
    }

	std::vector<TFeature> m_Features;
	std::queue<uint32_t> m_FeatureSlots;
    std::mutex m_Mutex;
};