//----------------------------------------------------------------------------------
// File:        CDx11Api.cpp
// SDK Version: 1.0.2
//
// SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: LicenseRef-NvidiaProprietary
//
// NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
// property and proprietary rights in and to this material, related
// documentation and any modifications thereto. Any use, reproduction,
// disclosure or distribution of this material and related documentation
// without an express license agreement from NVIDIA CORPORATION or
// its affiliates is strictly prohibited.
//
//----------------------------------------------------------------------------------

///////////////////////////////////////////
// CDx11Api.cpp
//
// Top level API into DX12 video devices
//

#include <initguid.h>
#include "CDx11Api.h"
#include "CDx1x_Common.h"

#pragma comment( lib, "d3d11" )


HRESULT CDx11Api::Init(CDx1xCommon* pCDx1xCommon, HWND hWndDisplay)
{
    HRESULT hr = S_OK;

    ID3D10Multithread* pMultiThread = nullptr;

    IDXGIAdapter4* pAdapter = pCDx1xCommon->GetDxgiAdapterForWnd(hWndDisplay);
    if (!pAdapter) return E_FAIL;

    UINT Flags = D3D11_CREATE_DEVICE_VIDEO_SUPPORT | D3D11_CREATE_DEVICE_BGRA_SUPPORT;
    D3D_FEATURE_LEVEL featureLevels[] = { D3D_FEATURE_LEVEL_11_1, D3D_FEATURE_LEVEL_11_0 };
    hr = D3D11CreateDevice(pAdapter,
                            D3D_DRIVER_TYPE_UNKNOWN,
                            nullptr,
                            Flags,
                            featureLevels, ARRAYSIZE(featureLevels),
                            D3D11_SDK_VERSION,
                            &m_pD3D11Device,
                            nullptr,
                            &m_pD3D11DeviceContext);
    SafeRelease(pAdapter);
    if (FAILED(hr)) return hr;

    hr = m_pD3D11Device->QueryInterface(__uuidof(ID3D11VideoDevice), (void**)&m_pDX11VideoDevice);
    if (FAILED(hr)) return hr;

    hr = m_pD3D11DeviceContext->QueryInterface(__uuidof(ID3D11VideoContext1), (void**)&m_pVideoContext);
    if (FAILED(hr)) return hr;

    // Need to explitly set the multithreaded mode for this device
    hr = m_pD3D11DeviceContext->QueryInterface(__uuidof(ID3D10Multithread), (void**)&pMultiThread);
    if (FAILED(hr)) return hr;

    pMultiThread->SetMultithreadProtected(TRUE);
    SafeRelease(pMultiThread);

    return hr;
}

void CDx11Api::ShutDown()
{
    SafeRelease(m_pD3D11Device);
    SafeRelease(m_pD3D11DeviceContext);
    SafeRelease(m_pDX11VideoDevice);
    SafeRelease(m_pVideoContext);
    SafeRelease(m_pVideoProcessorEnum);
    SafeRelease(m_pVideoProcessor);
}

HRESULT CDx11Api::CreateVideoProcessor(DXGI_FORMAT VPOutputFormat, UINT uInWidth, UINT uInHeight, UINT uOutWidth, UINT uOutHeight)
{
    HRESULT hr = S_OK;

    // this is just a hint. actual usage is done by SetOutputTargetRect and SetStreamDestRect(
    D3D11_VIDEO_PROCESSOR_CONTENT_DESC ContentDesc = {};
    ContentDesc.InputFrameFormat    = D3D11_VIDEO_FRAME_FORMAT_PROGRESSIVE;
    ContentDesc.InputWidth          = uInWidth;
    ContentDesc.InputHeight         = uInHeight;
    ContentDesc.OutputWidth         = uOutWidth;
    ContentDesc.OutputHeight        = uOutHeight;
    ContentDesc.Usage = D3D11_VIDEO_USAGE_PLAYBACK_NORMAL;

    hr = m_pDX11VideoDevice->CreateVideoProcessorEnumerator(&ContentDesc, &m_pVideoProcessorEnum);
    if (FAILED(hr)) return hr;

    UINT uiFlags = 0;
    hr = m_pVideoProcessorEnum->CheckVideoProcessorFormat(VPOutputFormat, &uiFlags);
    if (FAILED(hr) || 0 == (uiFlags & D3D11_VIDEO_PROCESSOR_FORMAT_SUPPORT_OUTPUT))
    {
        return E_INVALIDARG;
    }

    hr = m_pDX11VideoDevice->CreateVideoProcessor(m_pVideoProcessorEnum, 0, &m_pVideoProcessor);
    if (FAILED(hr)) return hr;


    // input format
    m_pVideoContext->VideoProcessorSetStreamFrameFormat(m_pVideoProcessor, 0, D3D11_VIDEO_FRAME_FORMAT_PROGRESSIVE);

    // Output rate (repeat frames)
    m_pVideoContext->VideoProcessorSetStreamOutputRate(m_pVideoProcessor, 0, D3D11_VIDEO_PROCESSOR_OUTPUT_RATE_NORMAL, TRUE, NULL);

    // Source rect
    // Set to FALSE to use the entire surface
    m_pVideoContext->VideoProcessorSetStreamSourceRect(m_pVideoProcessor, 0, FALSE, nullptr);

    RECT DstRect = { 0, 0, (LONG)uOutWidth, (LONG)uOutHeight };
    // output target rect
    m_pVideoContext->VideoProcessorSetOutputTargetRect(m_pVideoProcessor, TRUE, &DstRect);
    // Stream dest rect
    m_pVideoContext->VideoProcessorSetStreamDestRect(m_pVideoProcessor, 0, TRUE, &DstRect);

    // Output color space
    m_pVideoContext->VideoProcessorSetOutputColorSpace1(m_pVideoProcessor, DXGI_COLOR_SPACE_RGB_FULL_G22_NONE_P709);

    // Output background color (black)
    D3D11_VIDEO_COLOR backgroundColor = {};
    backgroundColor.RGBA.A = 1.0F;
    backgroundColor.RGBA.R = 1.0F * static_cast<float>(GetRValue(0)) / 255.0F;
    backgroundColor.RGBA.G = 1.0F * static_cast<float>(GetGValue(0)) / 255.0F;
    backgroundColor.RGBA.B = 1.0F * static_cast<float>(GetBValue(0)) / 255.0F;

    m_pVideoContext->VideoProcessorSetOutputBackgroundColor(m_pVideoProcessor, FALSE, &backgroundColor);

    return hr;
}

HRESULT CDx11Api::CreateSwapChain(HWND hWndDisplay, DXGI_FORMAT dispFormat, UINT uSwapWidth, UINT uSwapHeight, CDx11SwapChain** ppSwapChain)
{
    HRESULT hr = S_OK;

    IDXGIFactory7* pFactory = nullptr;

    hr = CreateDXGIFactory2(0, IID_PPV_ARGS(&pFactory));
    if (FAILED(hr)) return hr;
    
    // Get the DXGISwapChain1
    DXGI_SWAP_CHAIN_DESC1 scd = { 0 };
    scd.Width = uSwapWidth;
    scd.Height = uSwapHeight;
    scd.Format = dispFormat;
    scd.Stereo = false;
    scd.SampleDesc.Count = 1;
    scd.SampleDesc.Quality = 0;
    scd.BufferUsage = DXGI_USAGE_BACK_BUFFER | DXGI_USAGE_RENDER_TARGET_OUTPUT;
    scd.BufferCount = 8;
    scd.Scaling = DXGI_SCALING_STRETCH;
    scd.SwapEffect = DXGI_SWAP_EFFECT_FLIP_SEQUENTIAL;
    scd.Flags = 0;

    IDXGISwapChain1* pSwapChain = nullptr;
    IDXGISwapChain4* pSwapChain4 = nullptr;
    hr = pFactory->CreateSwapChainForHwnd(m_pD3D11Device, hWndDisplay, &scd, NULL, NULL, &pSwapChain);
    SafeRelease(pFactory);
    if (FAILED(hr)) return hr;
    hr = pSwapChain->SetFullscreenState(FALSE, NULL);
    if (FAILED(hr)) return hr;
    hr = pSwapChain->QueryInterface(__uuidof(IDXGISwapChain4), (LPVOID*)&pSwapChain4);
    if (FAILED(hr)) return hr;

    *ppSwapChain = new CDx11SwapChain(hWndDisplay, dispFormat, pSwapChain4);

    SafeRelease(pSwapChain);

    return hr;
}


HRESULT CDx11SwapChain::GetBuffer(ID3D11Texture2D** ppBuffer)
{
    return m_pSwapChain->GetBuffer(0, __uuidof(ID3D11Texture2D), (void**)ppBuffer);
}

HRESULT CDx11SwapChain::Present()
{
    DXGI_PRESENT_PARAMETERS DxgiPresentParams = { 0 };
    return m_pSwapChain->Present1(0, 0, &DxgiPresentParams);
}

DXGI_COLOR_SPACE_TYPE CDx11SwapChain::UpdateSwapChainColorSpace(CDx1xCommon* pCDx1xCommon, bool bSupportHDR)
{
    return pCDx1xCommon->UpdateSwapChainColorSpace(m_hWndDisplay, m_pSwapChain, m_DisplayFormat, bSupportHDR);
}

void CDx11Api::SetOutputColorSpace(DXGI_COLOR_SPACE_TYPE ColorSpace)
{
    m_pVideoContext->VideoProcessorSetOutputColorSpace1(m_pVideoProcessor, ColorSpace);
}

//-------------------------------------------------------------------

HRESULT CDx11Api::CreateTexture2D(  UINT                Width,
                                    UINT                Height,
                                    DXGI_FORMAT         dxgiFormat,
                                    UINT                NumBuffers,
                                    bool                bIsHwRT,
                                    bool                bIsCPUWrite,
                                    ID3D11Texture2D**   ppTexture
                                )
{
    D3D11_TEXTURE2D_DESC texture2d_desc = { 0 };
    texture2d_desc.Width = Width;
    texture2d_desc.Height = Height;
    texture2d_desc.MipLevels = 1;
    texture2d_desc.ArraySize = NumBuffers;
    texture2d_desc.SampleDesc.Count = 1;
    texture2d_desc.MiscFlags = 0;
    texture2d_desc.Format = dxgiFormat;
    if (bIsHwRT)
    {
        texture2d_desc.Usage = D3D11_USAGE_DEFAULT;
        // set UNORDERED_ACCESS (for kernel write) and SHADER_RESOURCE (for kernel read) for all textures
        texture2d_desc.BindFlags = D3D11_BIND_UNORDERED_ACCESS 
                                 | D3D11_BIND_SHADER_RESOURCE;
        if (   dxgiFormat == DXGI_FORMAT_NV12
            || dxgiFormat == DXGI_FORMAT_YUY2
            || dxgiFormat == DXGI_FORMAT_P010
            || dxgiFormat == DXGI_FORMAT_AYUV
            || dxgiFormat == DXGI_FORMAT_Y410
            || dxgiFormat == DXGI_FORMAT_Y416)
        {
            // set DECODER for all YUV textures
            texture2d_desc.BindFlags |= D3D11_BIND_DECODER;
        }
        else
        {
            // set RENDER_TARGET for all non-YUV textures.
            texture2d_desc.BindFlags |= D3D11_BIND_RENDER_TARGET;
        }
    }
    else if (bIsCPUWrite)
    {
        texture2d_desc.Usage = D3D11_USAGE_DYNAMIC;
        texture2d_desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
        texture2d_desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
    }
    else
    {
        texture2d_desc.Usage = D3D11_USAGE_STAGING;
        texture2d_desc.BindFlags = 0;
        texture2d_desc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
    }
    HRESULT hr = m_pD3D11Device->CreateTexture2D(&texture2d_desc, nullptr, ppTexture);

    return hr;
}

HRESULT CDx11Api::VpBltHd(ID3D11Texture2D** ppSrcSurface, ID3D11Texture2D* pDstSurface, UINT uOutWidth, UINT uOutHeight)
{
    HRESULT hr = S_OK;
    ID3D11VideoProcessorInputView* pInputView = nullptr;
    ID3D11VideoProcessorOutputView* pOutputView = nullptr;
    RECT DstRect = { 0, 0, (LONG)uOutWidth, (LONG)uOutHeight };
    // output target rect
    m_pVideoContext->VideoProcessorSetOutputTargetRect(m_pVideoProcessor, TRUE, &DstRect);
    // Stream dest rect
    m_pVideoContext->VideoProcessorSetStreamDestRect(m_pVideoProcessor, 0, TRUE, &DstRect);


    D3D11_VIDEO_PROCESSOR_OUTPUT_VIEW_DESC OutputViewDesc;
    ZeroMemory(&OutputViewDesc, sizeof(OutputViewDesc));
    OutputViewDesc.ViewDimension = D3D11_VPOV_DIMENSION_TEXTURE2D;
    OutputViewDesc.Texture2D.MipSlice = 0;
    OutputViewDesc.Texture2DArray.MipSlice = 0;
    OutputViewDesc.Texture2DArray.FirstArraySlice = 0;
    hr = m_pDX11VideoDevice->CreateVideoProcessorOutputView(pDstSurface, m_pVideoProcessorEnum, &OutputViewDesc, &pOutputView);
    if (FAILED(hr)) return hr;

    ID3D11Resource* pD3D11Resource = nullptr;
    hr = ppSrcSurface[0]->QueryInterface(IID_ID3D11Resource, (void**)&pD3D11Resource);
    if (FAILED(hr)) return hr;
    D3D11_VIDEO_PROCESSOR_INPUT_VIEW_DESC InputViewDesc = { 0 };
    InputViewDesc.ViewDimension = D3D11_VPIV_DIMENSION_TEXTURE2D;
    InputViewDesc.Texture2D.ArraySlice = 0;
    hr = m_pDX11VideoDevice->CreateVideoProcessorInputView(pD3D11Resource, m_pVideoProcessorEnum, &InputViewDesc, &pInputView);
    SafeRelease(pD3D11Resource);
    if (FAILED(hr)) return hr;

    D3D11_VIDEO_PROCESSOR_STREAM StreamData = {};
    StreamData.Enable = TRUE;
    StreamData.pInputSurface = pInputView;

    hr = m_pVideoContext->VideoProcessorBlt(m_pVideoProcessor, pOutputView, 0, 1, &StreamData);
    if (FAILED(hr)) return hr;

    SafeRelease(pInputView);
    SafeRelease(pOutputView);
    return hr;
}


HRESULT CDx11Api::MapStagingSurface(ID3D11Texture2D* pStagingSurface, D3D11_MAPPED_SUBRESOURCE* pMapped, bool bToWrite)
{
    // Map the surface to fill with data
    HRESULT hr = m_pD3D11DeviceContext->Map(pStagingSurface, 0, bToWrite ? D3D11_MAP_WRITE_DISCARD : D3D11_MAP_READ, 0, pMapped);
    return hr;
}

HRESULT CDx11Api::UnMapStagingSurface(ID3D11Texture2D* pSurface, ID3D11Texture2D* pStagingSurface, bool bDoCopy, UINT uSubResource)
{
    // unmap the surface
    m_pD3D11DeviceContext->Unmap(pStagingSurface, 0);

    if (bDoCopy)
    {
        // copy to target surface
        m_pD3D11DeviceContext->CopySubresourceRegion(pSurface, uSubResource, 0, 0, 0, pStagingSurface, 0, NULL);
        m_pD3D11DeviceContext->Flush();
    }
    return S_OK;
}
