#include "FSR2Upscaler_DX11.h"

#include "ffx-fsr2-api/dx11/ffx_fsr2_dx11.h"

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

    m_DX11Device = ffxGetDeviceDX11(device);

    size_t scratchBufferSize = ffxFsr2GetScratchMemorySizeDX11();
    void* scratchBuffer = malloc(scratchBufferSize);
    ffxFsr2GetInterfaceDX11(&m_Fsr2Interface, device, scratchBuffer, scratchBufferSize);
    return true;
}

bool FSR2Upscaler_DX11::IsValidFeature(FSR2Feature& feature)
{
    return feature.isValid;
}

bool FSR2Upscaler_DX11::InitFeature(FSR2Feature& feature, const FSR3CommandInitializationData* initData)
{
    FfxFsr2ContextDescription contextDescription{};
    contextDescription.callbacks = m_Fsr2Interface;
    contextDescription.device = m_DX11Device;
    contextDescription.displaySize = { initData->displaySizeWidth, initData->displaySizeHeight };
    contextDescription.maxRenderSize = { initData->maxRenderSizeWidth, initData->maxRenderSizeHeight };
    contextDescription.flags = initData->flags;

    feature.isValid = FFX_OK == ffxFsr2ContextCreate(&feature.fsr2Context, &contextDescription);
    return feature.isValid;
}

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

void FSR2Upscaler_DX11::Execute(FSR2Feature& feature, const FSR3CommandExecutionData* execData)
{
    FfxFsr2DispatchDescription dispatchDescription{};

    ID3D11DeviceContext* ctx = nullptr;
    m_GraphicsDevice->GetDevice()->GetImmediateContext(&ctx);
    dispatchDescription.commandList = ctx;

    dispatchDescription.color = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.colorInput.image);
    dispatchDescription.depth = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.depth.image);
    dispatchDescription.motionVectors = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.motionVectors.image);
    dispatchDescription.exposure = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.exposureTexture.image);
    dispatchDescription.reactive = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.reactiveMask.image);
    dispatchDescription.transparencyAndComposition = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.transparencyMask.image);
    dispatchDescription.output = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.colorOutput.image, nullptr, FFX_RESOURCE_STATE_UNORDERED_ACCESS);

    if (dispatchDescription.reactive.resource == nullptr)
        dispatchDescription.reactive = ffxGetResourceDX11(&feature.fsr2Context, (ID3D11Resource*)feature.textureTable.biasColorMask.image);

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

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

    ffxFsr2ContextDispatch(&feature.fsr2Context, &dispatchDescription);
}

void FSR2Upscaler_DX11::AwaitEndOfFrame(uint64_t frameValue)
{
    // Nothing to do
}

void FSR2Upscaler_DX11::DestroyContext(FSR2Feature& feature)
{
    ffxFsr2ContextDestroy(&feature.fsr2Context);
}

void FSR2Upscaler_DX11::DoShutdown()
{
    if (m_Fsr2Interface.scratchBuffer != nullptr)
    {
        free(m_Fsr2Interface.scratchBuffer);
        m_Fsr2Interface.scratchBuffer = nullptr;
    }

    m_DX11Device = nullptr;
}