#include "FSR3Upscaler_DX12.h"

#include <d3d12.h>
#include <dxgi.h>

bool FSR3Upscaler_DX12::Init()
{
    ID3D12Device* device = m_GraphicsDevice->GetDevice();
    if (device == nullptr)
        return false;

    m_DX12BackendDesc.device = device;
    m_FrameFenceEventHandle = CreateEventEx(nullptr, nullptr, 0, EVENT_ALL_ACCESS);

    return LoadFidelityFXLibrary(TEXT("amd_fidelityfx_dx12.dll"), device);
}

bool FSR3Upscaler_DX12::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_DX12BackendDesc.header), nullptr);
}

void FSR3Upscaler_DX12::SetTexture(FSR3Texture textureType, UnityTextureID textureID, uint32_t width, uint32_t height, FSR3TextureDesc* outTextureDesc)
{
    outTextureDesc->image = (intptr_t)m_GraphicsDevice->TextureFromNativeTexture(textureID);
}

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

    UnityGraphicsD3D12RecordingState state;
    m_GraphicsDevice->CommandRecordingState(&state);
    dispatchUpscale.commandList = state.commandList;

    // Keep track of which frame this dispatch is happening on, so we can verify when it's finished
    feature.dispatchFrameValue = m_GraphicsDevice->GetNextFrameFenceValue();

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

    if (dispatchUpscale.reactive.resource == nullptr)
        dispatchUpscale.reactive = ffxApiGetResourceDX12((ID3D12Resource*)feature.textureTable.biasColorMask.image, FFX_API_RESOURCE_STATE_PIXEL_COMPUTE_READ);

    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_DX12::AwaitEndOfFrame(uint64_t frameValue)
{
    // Thanks to https://alextardif.com/D3D11To12P1.html for explaining how fences work and how to make the CPU wait for a command queue to complete.
    ID3D12Fence* frameFence = m_GraphicsDevice->GetFrameFence();
    if (frameValue > frameFence->GetCompletedValue())
    {
        frameFence->SetEventOnCompletion(frameValue, m_FrameFenceEventHandle);
        WaitForSingleObjectEx(m_FrameFenceEventHandle, INFINITE, false);
    }
}

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

    if (m_FrameFenceEventHandle != nullptr)
    {
        CloseHandle(m_FrameFenceEventHandle);
        m_FrameFenceEventHandle = nullptr;
    }

    m_DX12BackendDesc.device = nullptr;
}