#include "NL_ImguiD3D11Renderer.h"

#include <d3d11.h>
#include <d3dcompiler.h>
#include <imgui.h>

namespace NLE
{
	namespace GRAPHICS
	{
		struct VERTEX_CONSTANT_BUFFER
		{
			float        mvp[4][4];
		};

		ImguiD3D11Renderer::ImguiD3D11Renderer() :
			_vertexBuffer(nullptr),
			_indexBuffer(nullptr),
			_vertexShaderBlob(nullptr),
			_vertexShader(nullptr),
			_inputLayout(nullptr),
			_vertexConstantBuffer(nullptr),
			_pixelShaderBlob(nullptr),
			_pixelShader(nullptr),
			_fontSampler(nullptr),
			_fontTextureView(nullptr),
			_rasterizerState(nullptr),
			_blendState(nullptr),
			_depthStencilState(nullptr),
			_vertexBufferSize(5000),
			_indexBufferSize(10000)
		{
		}

		ImguiD3D11Renderer::~ImguiD3D11Renderer()
		{

		}

		bool ImguiD3D11Renderer::initialize(ID3D11Device* device, ID3D11DeviceContext* devCon)
		{
			if (!createDeviceObjects(device, devCon))
				return false;

			if (!createFontAndTextures(device, devCon))
				return false;

			return true;
		}

		bool ImguiD3D11Renderer::createDeviceObjects(ID3D11Device* device, ID3D11DeviceContext* devCon)
		{
			// By using D3DCompile() from <d3dcompiler.h> / d3dcompiler.lib, we introduce a dependency to a given version of d3dcompiler_XX.dll (see D3DCOMPILER_DLL_A)
			// If you would like to use this DX11 sample code but remove this dependency you can: 
			//  1) compile once, save the compiled shader blobs into a file or source code and pass them to CreateVertexShader()/CreatePixelShader() [prefered solution]
			//  2) use code to detect any version of the DLL and grab a pointer to D3DCompile from the DLL. 
			// See https://github.com/ocornut/imgui/pull/638 for sources and details.

			// Create the vertex shader
			{
				static const char* vertexShader =
					"cbuffer vertexBuffer : register(b0) \
            {\
            float4x4 ProjectionMatrix; \
            };\
            struct VS_INPUT\
            {\
            float2 pos : POSITION;\
            float4 col : COLOR0;\
            float2 uv  : TEXCOORD0;\
            };\
            \
            struct PS_INPUT\
            {\
            float4 pos : SV_POSITION;\
            float4 col : COLOR0;\
            float2 uv  : TEXCOORD0;\
            };\
            \
            PS_INPUT main(VS_INPUT input)\
            {\
            PS_INPUT output;\
            output.pos = mul( ProjectionMatrix, float4(input.pos.xy, 0.f, 1.f));\
            output.col = input.col;\
            output.uv  = input.uv;\
            return output;\
            }";

				D3DCompile(vertexShader, strlen(vertexShader), nullptr, nullptr, nullptr, "main", "vs_4_0", 0, 0, &_vertexShaderBlob, nullptr);
				if (!_vertexShaderBlob) // NB: Pass ID3D10Blob* pErrorBlob to D3DCompile() to get error showing in (const char*)pErrorBlob->GetBufferPointer(). Make sure to Release() the blob!
					return false;

				if (device->CreateVertexShader((DWORD*)_vertexShaderBlob->GetBufferPointer(), _vertexShaderBlob->GetBufferSize(), nullptr, &_vertexShader) != S_OK)
					return false;

				// Create the input layout
				D3D11_INPUT_ELEMENT_DESC local_layout[] = {
					{ "POSITION", 0, DXGI_FORMAT_R32G32_FLOAT,   0, (size_t)(&((ImDrawVert*)0)->pos), D3D11_INPUT_PER_VERTEX_DATA, 0 },
					{ "TEXCOORD", 0, DXGI_FORMAT_R32G32_FLOAT,   0, (size_t)(&((ImDrawVert*)0)->uv),  D3D11_INPUT_PER_VERTEX_DATA, 0 },
					{ "COLOR",    0, DXGI_FORMAT_R8G8B8A8_UNORM, 0, (size_t)(&((ImDrawVert*)0)->col), D3D11_INPUT_PER_VERTEX_DATA, 0 },
				};

				if (device->CreateInputLayout(local_layout, 3, _vertexShaderBlob->GetBufferPointer(), _vertexShaderBlob->GetBufferSize(), &_inputLayout) != S_OK)
					return false;

				// Create the constant buffer
				{
					D3D11_BUFFER_DESC desc;
					desc.ByteWidth = sizeof(VERTEX_CONSTANT_BUFFER);
					desc.Usage = D3D11_USAGE_DYNAMIC;
					desc.BindFlags = D3D11_BIND_CONSTANT_BUFFER;
					desc.CPUAccessFlags = D3D11_CPU_ACCESS_WRITE;
					desc.MiscFlags = 0;
					device->CreateBuffer(&desc, nullptr, &_vertexConstantBuffer);
				}
			}

			// Create the pixel shader
			{
				static const char* pixelShader =
					"struct PS_INPUT\
            {\
            float4 pos : SV_POSITION;\
            float4 col : COLOR0;\
            float2 uv  : TEXCOORD0;\
            };\
            sampler sampler0;\
            Texture2D texture0;\
            \
            float4 main(PS_INPUT input) : SV_Target\
            {\
            float4 out_col = input.col * texture0.Sample(sampler0, input.uv); \
            return out_col; \
            }";

				D3DCompile(pixelShader, strlen(pixelShader), nullptr, nullptr, nullptr, "main", "ps_4_0", 0, 0, &_pixelShaderBlob, nullptr);
				if (!_pixelShaderBlob)  // NB: Pass ID3D10Blob* pErrorBlob to D3DCompile() to get error showing in (const char*)pErrorBlob->GetBufferPointer(). Make sure to Release() the blob!
					return false;

				if (device->CreatePixelShader((DWORD*)_pixelShaderBlob->GetBufferPointer(), _pixelShaderBlob->GetBufferSize(), nullptr, &_pixelShader) != S_OK)
					return false;
			}

			// Create the blending setup
			{
				D3D11_BLEND_DESC desc;
				ZeroMemory(&desc, sizeof(desc));
				desc.AlphaToCoverageEnable = false;
				desc.RenderTarget[0].BlendEnable = true;
				desc.RenderTarget[0].SrcBlend = D3D11_BLEND_SRC_ALPHA;
				desc.RenderTarget[0].DestBlend = D3D11_BLEND_INV_SRC_ALPHA;
				desc.RenderTarget[0].BlendOp = D3D11_BLEND_OP_ADD;
				desc.RenderTarget[0].SrcBlendAlpha = D3D11_BLEND_INV_SRC_ALPHA;
				desc.RenderTarget[0].DestBlendAlpha = D3D11_BLEND_ZERO;
				desc.RenderTarget[0].BlendOpAlpha = D3D11_BLEND_OP_ADD;
				desc.RenderTarget[0].RenderTargetWriteMask = D3D11_COLOR_WRITE_ENABLE_ALL;
				device->CreateBlendState(&desc, &_blendState);
			}

			// Create the rasterizer state
			{
				D3D11_RASTERIZER_DESC desc;
				ZeroMemory(&desc, sizeof(desc));
				desc.FillMode = D3D11_FILL_SOLID;
				desc.CullMode = D3D11_CULL_NONE;
				desc.ScissorEnable = true;
				desc.DepthClipEnable = true;
				device->CreateRasterizerState(&desc, &_rasterizerState);
			}

			// Create depth-stencil State
			{
				D3D11_DEPTH_STENCIL_DESC desc;
				ZeroMemory(&desc, sizeof(desc));
				desc.DepthEnable = false;
				desc.DepthWriteMask = D3D11_DEPTH_WRITE_MASK_ALL;
				desc.DepthFunc = D3D11_COMPARISON_ALWAYS;
				desc.StencilEnable = false;
				desc.FrontFace.StencilFailOp = desc.FrontFace.StencilDepthFailOp = desc.FrontFace.StencilPassOp = D3D11_STENCIL_OP_KEEP;
				desc.FrontFace.StencilFunc = D3D11_COMPARISON_ALWAYS;
				desc.BackFace = desc.FrontFace;
				device->CreateDepthStencilState(&desc, &_depthStencilState);
			}

			return true;
		}

		bool ImguiD3D11Renderer::createFontAndTextures(ID3D11Device* device, ID3D11DeviceContext* devCon)
		{
			// Build texture atlas
			ImGuiIO& io = ImGui::GetIO();
			unsigned char* pixels;
			int width, height;
			io.Fonts->GetTexDataAsRGBA32(&pixels, &width, &height);

			// Upload texture to graphics system
			{
				D3D11_TEXTURE2D_DESC desc;
				ZeroMemory(&desc, sizeof(desc));
				desc.Width = width;
				desc.Height = height;
				desc.MipLevels = 1;
				desc.ArraySize = 1;
				desc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
				desc.SampleDesc.Count = 1;
				desc.Usage = D3D11_USAGE_DEFAULT;
				desc.BindFlags = D3D11_BIND_SHADER_RESOURCE;
				desc.CPUAccessFlags = 0;

				ID