/**
    Copyright (c) 2010 yakiimo02
    Distributed under the New BSD License.
    See included license.txt or http://www.yakiimo3d.com/NewBSDLicense.txt
**/

#include "DXUT.h"
#include "CaptureFrameBuffer.h"
#include "DXHelper.h"

#include <string>

CaptureFrameBuffer    g_captureFrameBuffer;

/**
*/
CaptureFrameBuffer::CaptureFrameBuffer()
{
}

/**
    @param strFileName [in]  file name not file path.
*/
HRESULT CaptureFrameBuffer::Capture( const std::wstring& strFileName )
{
    HRESULT hr = S_OK;

    ID3D11DeviceContext* pd3dImmediateContext = DXUTGetD3D11DeviceContext();
    ID3D11Device* pd3dDevice = DXUTGetD3D11Device();

    // Save screenshot in same directory as executable.
    wchar_t buffer[ MAX_PATH ];
    int nLen = ::GetModuleFileName(NULL, buffer, MAX_PATH);
    if( nLen == 0 ) {
       return E_FAIL;
    }
    std::wstring strExeFullPath = buffer;

    std::wstring strOutFilePath = strExeFullPath.substr( 0, strExeFullPath.rfind( L"\\" )+1 ) + strFileName;

    // Retrieve RT resource
    ID3D11Resource *pRTResource;
    DXUTGetD3D11RenderTargetView()->GetResource(&pRTResource);

    // Retrieve a Texture2D interface from resource
    ID3D11Texture2D* RTTexture;
    pRTResource->QueryInterface( __uuidof( ID3D11Texture2D ), ( LPVOID* )&RTTexture);

    // We need a screen-sized STAGING resource for frame capturing
    const DXGI_SURFACE_DESC* pBackBufferSurfaceDesc = DXUTGetDXGIBackBufferSurfaceDesc();

    ID3D11Texture2D*  pCaptureTexture;
    {
        D3D11_TEXTURE2D_DESC TexDesc;
        DXGI_SAMPLE_DESC SingleSample = { 1, 0 };
        TexDesc.Width = pBackBufferSurfaceDesc->Width;
        TexDesc.Height = pBackBufferSurfaceDesc->Height;
        TexDesc.Format = pBackBufferSurfaceDesc->Format;
        TexDesc.SampleDesc = SingleSample;
        TexDesc.MipLevels = 1;
        TexDesc.Usage = D3D11_USAGE_STAGING;
        TexDesc.MiscFlags = 0;
        TexDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        TexDesc.BindFlags = 0;
        TexDesc.ArraySize = 1;
        V_RETURN( pd3dDevice->CreateTexture2D( &TexDesc, NULL, &pCaptureTexture ) );
    }

    // Check if RT is multisampled or not
    D3D11_TEXTURE2D_DESC    TexDesc;
    RTTexture->GetDesc(&TexDesc);
    if (TexDesc.SampleDesc.Count>1)
    {
        // RT is multisampled, need resolving before dumping to disk
        // Create single-sample RT of the same type and dimensions
        DXGI_SAMPLE_DESC SingleSample = { 1, 0 };
        TexDesc.CPUAccessFlags = D3D11_CPU_ACCESS_READ;
        TexDesc.MipLevels = 1;
        TexDesc.Usage = D3D11_USAGE_DEFAULT;
        TexDesc.CPUAccessFlags = 0;
        TexDesc.BindFlags = 0;
        TexDesc.SampleDesc = SingleSample;

        ID3D11Texture2D *pSingleSampleTexture;
        DXUTGetD3D11Device()->CreateTexture2D(&TexDesc, NULL, &pSingleSampleTexture);

        DXUTGetD3D11DeviceContext()->ResolveSubresource(pSingleSampleTexture, 0, RTTexture, 0, TexDesc.Format);

        // Copy RT into STAGING texture
        DXUTGetD3D11DeviceContext()->CopyResource(pCaptureTexture, pSingleSampleTexture);

        D3DX11SaveTextureToFile(pd3dImmediateContext, pCaptureTexture, D3DX11_IFF_BMP, strOutFilePath.c_str());

        SAFE_RELEASE(pSingleSampleTexture);
        
    }
    else
    {
        // Single sample case
        // Copy RT into STAGING texture
        DXUTGetD3D11DeviceContext()->CopyResource(pCaptureTexture, pRTResource);

        D3DX11SaveTextureToFile(pd3dImmediateContext, pCaptureTexture, D3DX11_IFF_BMP, strOutFilePath.c_str() );
    }

    SAFE_RELEASE(RTTexture);
    SAFE_RELEASE(pRTResource);
    SAFE_RELEASE(pCaptureTexture);

    return hr;
}
