#include "orcusExportShaderOutput.h"
#include <ScreenGrab.h>

#define SAFE_RELEASE(x) do { if (x) { (x)->Release(); x = NULL; } } while(false)

orcusExport::orcusExport()
	: device(NULL),
	devCon(NULL),
	texRT(NULL),
	rt(NULL),
	rtOld(NULL),
	dsOld(NULL),
	bsOld(NULL)
{
}

void orcusExport::prePS(ID3D11Device * device, ID3D11DeviceContext * devCon)
{
	this->device = device;
	this->devCon = devCon;

	D3D11_VIEWPORT vp;
	UINT vpNum = 1;
	devCon->RSGetViewports(&vpNum, &vp);

	D3D11_TEXTURE2D_DESC td;
	ZeroMemory(&td, sizeof(td));
	td.Width = vp.Width;
	td.Height = vp.Height;
	td.MipLevels = 1;
	td.ArraySize = 1;
	td.Format = DXGI_FORMAT_R32G32B32A32_FLOAT;
	td.SampleDesc.Count = 1;
	td.Usage = D3D11_USAGE_DEFAULT;
	td.BindFlags = D3D11_BIND_RENDER_TARGET | D3D11_BIND_SHADER_RESOURCE;
	td.CPUAccessFlags = 0;
	td.MiscFlags = 0;
	device->CreateTexture2D(&td, NULL, &texRT);

	D3D11_RENDER_TARGET_VIEW_DESC rtd;
	rtd.Format = td.Format;
	rtd.ViewDimension = D3D11_RTV_DIMENSION_TEXTURE2D;
	rtd.Texture2D.MipSlice = 0;
	device->CreateRenderTargetView(texRT, &rtd, &rt);

	devCon->OMGetRenderTargets(1, &rtOld, &dsOld);
	devCon->OMSetRenderTargets(1, &rt, dsOld);

	devCon->OMGetBlendState(&bsOld, blendFactor, &sampleMask);
	D3D11_BLEND_DESC bsd;
	bsd.AlphaToCoverageEnable = FALSE;
	bsd.IndependentBlendEnable = FALSE;
	bsd.RenderTarget[0].BlendEnable = TRUE;
	bsd.RenderTarget[0].BlendOp = D3D11_BLEND_OP_ADD;
	bsd.RenderTarget[0].SrcBlend = D3D11_BLEND_ONE;
	bsd.RenderTarget[0].DestBlend = D3D11_BLEND_ZERO;
	bsd.RenderTarget[0].BlendOpAlpha = D3D11_BLEND_OP_ADD;
	bsd.RenderTarget[0].SrcBlendAlpha = D3D11_BLEND_ONE;
	bsd.RenderTarget[0].DestBlendAlpha = D3D11_BLEND_ZERO;
	bsd.RenderTarget[0].RenderTargetWriteMask = D3D11_COLOR_WRITE_ENABLE_ALL;
	device->CreateBlendState(&bsd, &bs);
	devCon->OMSetBlendState(bs, blendFactor, sampleMask);

	float color[4] = { 0, 0, 0, 0 };
	devCon->ClearRenderTargetView(rt, color);
}

void orcusExport::postPS(const std::string & texFilename)
{
	int wlen = MultiByteToWideChar(CP_ACP, 0, texFilename.c_str(), -1, NULL, 0);
	wchar_t * wFilename = new wchar_t[wlen];
	MultiByteToWideChar(CP_ACP, 0, texFilename.c_str(), -1, wFilename, wlen);
	DWORD attrib = GetFileAttributes(wFilename);
	if (attrib == INVALID_FILE_ATTRIBUTES)
		DirectX::SaveDDSTextureToFile(devCon, texRT, wFilename);
	delete[] wFilename;
	devCon->OMSetBlendState(bsOld, blendFactor, sampleMask);
	devCon->OMSetRenderTargets(1, &rtOld, dsOld);
	SAFE_RELEASE(bs);
	SAFE_RELEASE(bsOld);
	SAFE_RELEASE(rt);
	SAFE_RELEASE(rtOld);
	SAFE_RELEASE(dsOld);
	SAFE_RELEASE(texRT);
}