#include "FSR3Upscaler_Vulkan.h"

static VkInstance s_VulkanInstance;
static PFN_vkGetInstanceProcAddr s_GetInstanceProcAddr;

static PFN_vkVoidFunction GetVulkanDeviceProcAddr(VkDevice device, const char* pName)
{
    if (s_GetInstanceProcAddr == nullptr)
        return nullptr;

    return s_GetInstanceProcAddr(s_VulkanInstance, pName);
}

bool FSR3Upscaler_Vulkan::Init()
{
    UnityVulkanInstance instance = m_GraphicsDevice->Instance();
    if (instance.device == nullptr)
        return false;

    s_VulkanInstance = instance.instance;
    s_GetInstanceProcAddr = instance.getInstanceProcAddr;

    m_VulkanBackendDesc.vkDevice = instance.device;
    m_VulkanBackendDesc.vkPhysicalDevice = instance.physicalDevice;
    m_VulkanBackendDesc.vkDeviceProcAddr = &GetVulkanDeviceProcAddr;

    return LoadFidelityFXLibrary(TEXT("amd_fidelityfx_vk.dll"), nullptr);
}

bool FSR3Upscaler_Vulkan::InitFeature(FSR3Feature_FFX& feature, const FSR3CommandInitializationData* initData)
{
    ffx::CreateContextDescUpscale createUpscaling;
    createUpscaling.maxUpscaleSize = { initData->displaySizeWidth, initData->displaySizeHeight };
    createUpscaling.maxRenderSize = { initData->maxRenderSizeWidth, initData->maxRenderSizeHeight };
    createUpscaling.flags = initData->flags;

	return FFX_API_RETURN_OK == m_ffxFunctions.CreateContext(&feature.upscalingContext, ffx::LinkHeaders(createUpscaling.header, m_VulkanBackendDesc.header), nullptr);
}

void FSR3Upscaler_Vulkan::SetTexture(FSR3Texture textureType, UnityTextureID textureID, uint32_t width, uint32_t height, FSR3TextureDesc* outTextureDesc)
{
    VkAccessFlags accessFlags = (textureType == FSR3Texture::tColorOutput) ? VK_ACCESS_SHADER_WRITE_BIT : VK_ACCESS_SHADER_READ_BIT;

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

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);
}

void FSR3Upscaler_Vulkan::Execute(FSR3Feature_FFX& feature, const FSR3CommandExecutionData* execData)
{
    ffx::DispatchDescUpscale dispatchUpscale{};

    UnityVulkanRecordingState state;
    m_GraphicsDevice->CommandRecordingState(&state, kUnityVulkanGraphicsQueueAccess_DontCare);
    dispatchUpscale.commandList = state.commandBuffer;

    // Keep track of which frame this dispatch is happening on, so we can verify when it's finished
    feature.dispatchFrameValue = state.currentFrameNumber;

    UnityVulkanInstance instance = m_GraphicsDevice->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);

    if (dispatchUpscale.reactive.resource == nullptr)
        dispatchUpscale.reactive = GetVulkanTextureResource(instance, feature.textureTable.biasColorMask);

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

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

    m_ffxFunctions.Dispatch(&feature.upscalingContext, &dispatchUpscale.header);
}

void FSR3Upscaler_Vulkan::AwaitEndOfFrame(uint64_t frameValue)
{
    UnityVulkanRecordingState state;
    m_GraphicsDevice->CommandRecordingState(&state, kUnityVulkanGraphicsQueueAccess_DontCare);
    if (frameValue > state.safeFrameNumber)
    {
        UnityVulkanInstance instance = m_GraphicsDevice->Instance();
        vkQueueWaitIdle(instance.graphicsQueue);
    }
}

void FSR3Upscaler_Vulkan::DoShutdown()
{
    FSR3Upscaler_FFXBase::DoShutdown();

    m_VulkanBackendDesc.vkDevice = nullptr;
    m_VulkanBackendDesc.vkPhysicalDevice = nullptr;
    m_VulkanBackendDesc.vkDeviceProcAddr = nullptr;
}