#include <stdint.h>

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

#include "UnityPluginAPI/IUnityInterface.h"
#include "UnityPluginAPI/IUnityLog.h"
#include "UnityPluginAPI/IUnityGraphics.h"
#include "UnityPluginAPI/IUnityRenderingExtensions.h"

#include <d3d12.h>
#include <dxgi.h>
#include "UnityPluginAPI/IUnityGraphicsD3D12.h"
#include "UnityPluginAPI/IUnityGraphicsVulkan.h"

#include "ffx_api/ffx_api_loader.h"
#include "ffx_api/ffx_upscale.hpp"
#include "ffx_api/dx12/ffx_api_dx12.hpp"
#include "ffx_api/vk/ffx_api_vk.hpp"

#include "FSR3UnityTypes.h"

static const int32_t BaseEventId = 0;

static IUnityInterfaces* s_UnityInterfaces = nullptr;
static IUnityLog* s_Log = nullptr;
static IUnityGraphics* s_Graphics = nullptr;
static IUnityGraphicsD3D12v7* s_GraphicsD3D12 = nullptr;
static IUnityGraphicsVulkan* s_GraphicsVulkan = nullptr;
static UnityGfxRenderer s_RendererType = kUnityGfxRendererNull;

static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType);
static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data);
static void UNITY_INTERFACE_API OnSetTextureEvent(int eventID, void* data);

struct FSR3Feature
{
    ffx::Context upscalingContext;

    uint32_t upscaleSizeWidth;
    uint32_t upscaleSizeHeight;
    uint32_t flags;

    FSR3TextureTable textureTable;
};

static std::vector<FSR3Feature> s_Features;
static std::queue<uint32_t> s_FeatureSlots;
static std::mutex s_FeatureMutex;

static HMODULE s_ffxModule = nullptr;
static ffxFunctions s_ffxFunctions{};
static ffx::CreateBackendDX12Desc s_DX12BackendDesc{};
static ffx::CreateBackendVKDesc s_VulkanBackendDesc{};

// Unity plugin load event
extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginLoad(IUnityInterfaces* unityInterfaces)
{
    s_UnityInterfaces = unityInterfaces;
    s_Log = unityInterfaces->Get<IUnityLog>();
    s_Graphics = unityInterfaces->Get<IUnityGraphics>();

    s_Graphics->RegisterDeviceEventCallback(OnGraphicsDeviceEvent);

    // Run OnGraphicsDeviceEvent(initialize) manually on plugin load
    // to not miss the event in case the graphics device is already initialized
    OnGraphicsDeviceEvent(kUnityGfxDeviceEventInitialize);
}

// Unity plugin unload event
extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API UnityPluginUnload()
{
    s_Graphics->UnregisterDeviceEventCallback(OnGraphicsDeviceEvent);
}

static void UNITY_INTERFACE_API OnGraphicsDeviceEvent(UnityGfxDeviceEventType eventType)
{
    switch (eventType)
    {
    case kUnityGfxDeviceEventInitialize:
    {
        s_RendererType = s_Graphics->GetRenderer();
        switch (s_RendererType)
        {
        case kUnityGfxRendererD3D12:
        {
            s_GraphicsD3D12 = s_UnityInterfaces->Get<IUnityGraphicsD3D12v7>();
            if (s_GraphicsD3D12 == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain D3D12 Graphics interface!");
                return;
            }
            break;
        }
        case kUnityGfxRendererVulkan:
        {
            s_GraphicsVulkan = s_UnityInterfaces->Get<IUnityGraphicsVulkan>();
            if (s_GraphicsVulkan == nullptr)
            {
                UNITY_LOG_ERROR(s_Log, "Could not obtain Vulkan Graphics interface!");
                return;
            }
            break;
        }
        }

        break;
    }
    case kUnityGfxDeviceEventShutdown:
    {
        s_GraphicsVulkan = nullptr;
        s_GraphicsD3D12 = nullptr;
        s_RendererType = kUnityGfxRendererNull;
        break;
    }
    case kUnityGfxDeviceEventBeforeReset:
    {
        break;
    }
    case kUnityGfxDeviceEventAfterReset:
    {
        break;
    }
    };
}

// Thread-safe allocation of feature slot
static uint32_t AllocateFeatureSlot()
{
    std::lock_guard<std::mutex> lock(s_FeatureMutex);

    if (s_FeatureSlots.empty())
    {
        // Create a new feature if there are no free slots
        uint32_t featureSlot = (uint32_t)s_Features.size();
        s_Features.push_back(std::move(FSR3Feature()));
        return featureSlot;
    }

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

// Thread-safe freeing and clearing of feature slot
static void FreeFeatureSlot(uint32_t featureSlot)
{
    std::lock_guard<std::mutex> lock(s_FeatureMutex);

    s_FeatureSlots.push(featureSlot);
    memset(&s_Features[featureSlot], 0, sizeof(FSR3Feature));
}

static PFN_vkVoidFunction GetVulkanDeviceProcAddr(VkDevice device, const char* pName)
{
    UnityVulkanInstance instance = s_GraphicsVulkan->Instance();
    return instance.getInstanceProcAddr(instance.instance, pName);
}

static bool LoadFidelityFXLibrary(_In_ LPCWSTR lpLibFileName)
{
    s_ffxModule = LoadLibrary(lpLibFileName);
    if (s_ffxModule == nullptr)
    {
        UNITY_LOG_ERROR(s_Log, "Failed to load FidelityFX library!");
        return false;
    }

    ffxLoadFunctions(&s_ffxFunctions, s_ffxModule);
    if (s_ffxFunctions.CreateContext == nullptr)
    {
        UNITY_LOG_ERROR(s_Log, "Failed to load FidelityFX library functions!");
        return false;
    }

    return true;
}

extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_InitApi()
{
    if (s_GraphicsD3D12 != nullptr)
    {
        ID3D12Device* device = s_GraphicsD3D12->GetDevice();
        if (device == nullptr)
            return false;

        s_DX12BackendDesc.device = device;

        return LoadFidelityFXLibrary(TEXT("amd_fidelityfx_dx12.dll"));
    }
    else if (s_GraphicsVulkan != nullptr)
    {
        UnityVulkanInstance instance = s_GraphicsVulkan->Instance();
        if (instance.device == nullptr)
            return false;

        s_VulkanBackendDesc.vkDevice = instance.device;
        s_VulkanBackendDesc.vkPhysicalDevice = instance.physicalDevice;
        s_VulkanBackendDesc.vkDeviceProcAddr = &GetVulkanDeviceProcAddr;
        
        return LoadFidelityFXLibrary(TEXT("amd_fidelityfx_vk.dll"));
    }

    return false;
}

extern "C" void UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_ShutdownApi()
{
    for (uint32_t slot = 0; slot < s_Features.size(); ++slot)
    {
        auto& feature = s_Features[slot];
        if (feature.upscalingContext != nullptr)
        {
            s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
            FreeFeatureSlot(slot);
        }
    }

    if (s_ffxModule != nullptr)
    {
        FreeLibrary(s_ffxModule);
        s_ffxModule = nullptr;
    }
    
    s_VulkanBackendDesc.vkDevice = nullptr;
    s_VulkanBackendDesc.vkPhysicalDevice = nullptr;
    s_VulkanBackendDesc.vkDeviceProcAddr = nullptr;

    s_DX12BackendDesc.device = nullptr;
}

extern "C" uint32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetDeviceVersion()
{
    return 0x0u;
}

extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetRenderEventCallback()
{
    return OnRenderEventAndData;
}

extern "C" UnityRenderingEventAndData UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetSetTextureEventCallback()
{
    return OnSetTextureEvent;
}

extern "C" uint32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_CreateFeatureSlot()
{
    return AllocateFeatureSlot();
}

extern "C" float UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetUpscaleRatioFromQualityMode(FSR3Quality qualityMode)
{
    switch (qualityMode)
    {
    case FSR3Quality::qQuality:
        return 1.5f;
    case FSR3Quality::qBalanced:
        return 1.7f;
    case FSR3Quality::qPerformance:
        return 2.0f;
    case FSR3Quality::qUltraPerformance:
        return 3.0f;
    default:
        return 1.0f;
    }
}

extern "C" bool UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetRenderResolutionFromQualityMode(FSR3Quality qualityMode, uint32_t displayWidth, uint32_t displayHeight, uint32_t* renderWidth, uint32_t* renderHeight)
{
    if (renderWidth == nullptr || renderHeight == nullptr)
        return false;

    float ratio = AMDUP_GetUpscaleRatioFromQualityMode(qualityMode);
    *renderWidth = (uint32_t)roundf(displayWidth / ratio);
    *renderHeight = (uint32_t)roundf(displayHeight / ratio);
    return true;
}

extern "C" int32_t UNITY_INTERFACE_EXPORT UNITY_INTERFACE_API AMDUP_GetBaseEventId()
{
    return BaseEventId;
}

static FfxApiResource GetVulkanTextureResource(UnityVulkanInstance& instance, FSR3TextureDesc& texture, FfxApiResourceState state = FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ)
{
    FfxApiResourceDescription desc{};

    if (texture.image != 0)
    {
        desc.type = FFX_API_RESOURCE_TYPE_TEXTURE2D;
        desc.format = ffxApiGetSurfaceFormatVK((VkFormat)texture.format);
        desc.width = texture.width;
        desc.height = texture.height;
        desc.depth = 1;
        desc.mipCount = 1;
        desc.flags = FFX_API_RESOURCE_FLAGS_NONE;
        desc.usage = state == FFX_API_RESOURCE_STATE_UNORDERED_ACCESS ? FFX_API_RESOURCE_USAGE_UAV : FFX_API_RESOURCE_USAGE_READ_ONLY;
    }

    return ffxApiGetResourceVK((void*)texture.image, desc, state);
}

// Plugin function to handle a specific rendering event
static void UNITY_INTERFACE_API OnRenderEventAndData(int eventID, void* data)
{
    // User rendering code
    switch (eventID)
    {
    case BaseEventId + FSR3PluginEvent::eDestroyFeature:
    {
        uint32_t featureSlot = (uint32_t)(int64_t)data;
        if (featureSlot < 0 || featureSlot >= s_Features.size())
            return;

        auto& feature = s_Features[featureSlot];
        if (feature.upscalingContext != nullptr)
        {
            s_ffxFunctions.DestroyContext(&feature.upscalingContext, nullptr);
            FreeFeatureSlot(featureSlot);
        }
        break;
    }
    case BaseEventId + FSR3PluginEvent::eExecute:
    {
        auto* params = (FSR3CommandExecutionData*)data;
        if (params->featureSlot < 0 || params->featureSlot >= s_Features.size())
            return;

        auto& feature = s_Features[params->featureSlot];

        ffx::DispatchDescUpscale dispatchUpscale{};

        if (s_GraphicsD3D12 != nullptr)
        {
            UnityGraphicsD3D12RecordingState state;
            s_GraphicsD3D12->CommandRecordingState(&state);
            dispatchUpscale.commandList = state.commandList;

            dispatchUpscale.color = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.colorInput.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.depth = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.depth.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.motionVectors = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.motionVectors.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.exposure = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.exposureTexture.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.reactive = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.reactiveMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.transparencyAndComposition = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.transparencyMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);
            dispatchUpscale.output = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.colorOutput.image, FFX_API_RESOURCE_STATE_UNORDERED_ACCESS);
        }
        else if (s_GraphicsVulkan != nullptr)
        {
            UnityVulkanRecordingState state;
            s_GraphicsVulkan->CommandRecordingState(&state, kUnityVulkanGraphicsQueueAccess_DontCare);
            dispatchUpscale.commandList = state.commandBuffer;

            UnityVulkanInstance instance = s_GraphicsVulkan->Instance();
            dispatchUpscale.color = GetVulkanTextureResource(instance, feature.textureTable.colorInput);
            dispatchUpscale.depth = GetVulkanTextureResource(instance, feature.textureTable.depth);
            dispatchUpscale.motionVectors = GetVulkanTextureResource(instance, feature.textureTable.motionVectors);
            dispatchUpscale.exposure = GetVulkanTextureResource(instance, feature.textureTable.exposureTexture);
            dispatchUpscale.reactive = GetVulkanTextureResource(instance, feature.textureTable.reactiveMask);
            dispatchUpscale.transparencyAndComposition = GetVulkanTextureResource(instance, feature.textureTable.transparencyMask);
            dispatchUpscale.output = GetVulkanTextureResource(instance, feature.textureTable.colorOutput, FFX_API_RESOURCE_STATE_UNORDERED_ACCESS);
        }

        dispatchUpscale.jitterOffset.x = params->jitterOffsetX;
        dispatchUpscale.jitterOffset.y = params->jitterOffsetY;
        dispatchUpscale.motionVectorScale.x = params->MVScaleX;
        dispatchUpscale.motionVectorScale.y = params->MVScaleY;
        dispatchUpscale.reset = params->reset;
        dispatchUpscale.enableSharpening = params->enableSharpening;
        dispatchUpscale.sharpness = params->sharpness;
        dispatchUpscale.frameTimeDelta = params->frameTimeDelta;
        dispatchUpscale.preExposure = params->preExposure;
        dispatchUpscale.renderSize.width = params->renderSizeWidth;
        dispatchUpscale.renderSize.height = params->renderSizeHeight;
        dispatchUpscale.upscaleSize.width = feature.upscaleSizeWidth;
        dispatchUpscale.upscaleSize.height = feature.upscaleSizeHeight;
        dispatchUpscale.cameraFovAngleVertical = params->cameraFovAngleVertical;

        if (feature.flags & FFX_UPSCALE_ENABLE_DEPTH_INVERTED)
        {
            dispatchUpscale.cameraFar = params->cameraNear;
            dispatchUpscale.cameraNear = params->cameraFar;
        }
        else
        {
            dispatchUpscale.cameraFar = params->cameraFar;
            dispatchUpscale.cameraNear = params->cameraNear;
        }

        s_ffxFunctions.Dispatch(&feature.upscalingContext, &dispatchUpscale.header);
        break;
    }
    case BaseEventId + FSR3PluginEvent::ePostExecute:
    {
        auto* params = (FSR3CommandExecutionData*)data;
        break;
    }
    case BaseEventId + FSR3PluginEvent::eInit:
    {
        auto* params = (FSR3CommandInitializationData*)data;
        if (params->featureSlot < 0 || params->featureSlot >= s_Features.size())
            return;

        auto& feature = s_Features[params->featureSlot];
        feature.upscaleSizeWidth = params->displaySizeWidth;
        feature.upscaleSizeHeight = params->displaySizeHeight;
        feature.flags = params->flags;

        ffx::CreateContextDescUpscale createUpscaling;
        createUpscaling.maxUpscaleSize = { params->displaySizeWidth, params->displaySizeHeight };
        createUpscaling.maxRenderSize = { params->maxRenderSizeWidth, params->maxRenderSizeHeight };
        createUpscaling.flags = params->flags;

        if (s_GraphicsD3D12 != nullptr)
        {
            s_ffxFunctions.CreateContext(&feature.upscalingContext, ffx::LinkHeaders(createUpscaling.header, s_DX12BackendDesc.header), nullptr);
        }
        else if (s_GraphicsVulkan != nullptr)
        {
            s_ffxFunctions.CreateContext(&feature.upscalingContext, ffx::LinkHeaders(createUpscaling.header, s_VulkanBackendDesc.header), nullptr);
        }
        break;
    }
    }
}

static void UNITY_INTERFACE_API OnSetTextureEvent(int eventID, void* data)
{
    auto* params = (UnityRenderingExtTextureUpdateParamsV2*)data;

    // userData = (featureId & (int) ushort.MaxValue) << 16 | (textureSlot & (int) short.MaxValue) << 1 | (clearTextureTable ? 1 : 0);
    uint32_t featureSlot = (params->userData >> 16) & 0xFFFF;
    uint32_t textureSlot = (params->userData >> 1) & 0x7FFF;
    uint32_t clearTextureTable = params->userData & 0x1;

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

    auto& feature = s_Features[featureSlot];

    // User rendering code
    switch (eventID)
    {
    case kUnityRenderingExtEventUpdateTextureBeginV2:
    {
        if (clearTextureTable)
        {
            memset(&feature.textureTable, 0, sizeof(FSR3TextureTable));
        }
        break;
    }
    case kUnityRenderingExtEventUpdateTextureEndV2:
    {
        // 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*)&feature.textureTable) + textureSlot;

        if (s_GraphicsD3D12 != nullptr)
        {
            textureDesc->image = (intptr_t)s_GraphicsD3D12->TextureFromNativeTexture((UnityTextureID)params->textureID);
        }
        else if (s_GraphicsVulkan != nullptr)
        {
            VkAccessFlags accessFlags = (textureSlot == FSR3Textures::tColorOutput) ? VK_ACCESS_SHADER_WRITE_BIT : VK_ACCESS_SHADER_READ_BIT;

            UnityVulkanImage image;
            if (s_GraphicsVulkan->AccessTextureByID((UnityTextureID)params->textureID, UnityVulkanWholeImage,
                VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT, accessFlags,
                kUnityVulkanResourceAccess_PipelineBarrier, &image))
            {
                textureDesc->image = (intptr_t)image.image;
                textureDesc->width = params->width;
                textureDesc->height = params->height;
                textureDesc->format = (uint32_t)image.format;
            }
            else
            {
                textureDesc->image = 0;
            }
        }

        break;
    }
    }
}