#include "stdafx.h"
#include "Shader.h"
#include "Texture.h"
#include "Core/Convert.h"
#include "Exception.h"

#ifdef GUI_WIN32

namespace gui {

	// Compile a shader. Throws on error.
	static ComPtr<ID3D10Blob> compileShader(Shader *s, const char *model) {
		// Docs says to use D3D10X (it gives us a slightly newer compiler). However, we want to
		// avoid depending on d3d10x.dll, so we use the one from the OS. We can work around bugs
		// anyway.
		Str *code = s->source(shaderHLSL);
		const char *utf8code = code->utf8_str();
		UINT flags = D3D10_SHADER_OPTIMIZATION_LEVEL2; // We have 0 to 3.
#ifdef DEBUG
		// Good to have to debug our codegen.
		flags |= D3D10_SHADER_WARNINGS_ARE_ERRORS;
#endif
		ComPtr<ID3D10Blob> shaderOut;
		ComPtr<ID3D10Blob> errorOut;
		HRESULT r = D3D10CompileShader(utf8code, strlen(utf8code), NULL,
									NULL, NULL,
									"main", model, flags,
									&shaderOut.v, &errorOut.v);
		if (FAILED(r)) {
			// Get text from the error blob.
			if (errorOut.v) {
				const char *ptr = (const char *)errorOut->GetBufferPointer();
				size_t size = errorOut->GetBufferSize();

				Str *errors = new (s) Str(toWChar(s->engine(), ptr, size));
				throw new (s) GuiError(TO_S(s, S("Failed to compile the shader:\n") << code << S("\n") << errors));
			} else {
				throw new (s) GuiError(TO_S(s, S("Failed to compile the shader.") << ::toS(r).c_str()));
			}
		}

		return shaderOut;
	}


	static DXGI_FORMAT lookupFormat(Bool integer, Nat components) {
		// TODO: UINT or SINT?
		if (integer && components == 1) {
			return DXGI_FORMAT_R32_UINT;
		} else if (integer && components == 2) {
			return DXGI_FORMAT_R32G32_UINT;
		} else if (integer && components == 3) {
			return DXGI_FORMAT_R32G32B32_UINT;
		} else if (integer && components >= 4) {
			return DXGI_FORMAT_R32G32B32A32_UINT;
		} else if (!integer && components == 1) {
			return DXGI_FORMAT_R32_FLOAT;
		} else if (!integer && components == 2) {
			return DXGI_FORMAT_R32G32_FLOAT;
		} else if (!integer && components == 3) {
			return DXGI_FORMAT_R32G32B32_FLOAT;
		} else if (!integer && components >= 4) {
			return DXGI_FORMAT_R32G32B32A32_FLOAT;
		} else {
			return DXGI_FORMAT_UNKNOWN;
		}
	}

	D3DVertexShader::D3DVertexShader(VertexShader *shader, D3DSurface &surface) {
		ComPtr<ID3D10Blob> code = compileShader(shader, "vs_4_0");
		const void *codePtr = code->GetBufferPointer();
		size_t codeSize = code->GetBufferSize();

		HRESULT r = surface.dev()->CreateVertexShader(codePtr, codeSize, &this->shader.v);
		verify(shader->engine(), r, S("Failed to create a vertex shader: "));

		VertexDesc *desc = shader->inputDescription();

		// Note: We want to keep the DESCs stack allocated since the strings will be GC:d.
		const UINT maxInput = 50;
		D3D10_INPUT_ELEMENT_DESC input[maxInput];
		UINT inputCount = 0;
		for (Nat i = 0; i < desc->count(); i++) {
			const VertexDesc::Elem &e = desc->at(i);

			for (Nat component = 0; component < e.components; component += 4) {
				if (inputCount >= maxInput)
					throw new (shader) GuiError(TO_S(shader, S("Too many shader components!")));

				StrBuf *semName = new (shader) StrBuf();
				*semName << S("ID") << e.identifier << S("_");

				input[inputCount].SemanticName = semName->toS()->utf8_str();
				input[inputCount].SemanticIndex = component;
				input[inputCount].Format = lookupFormat(e.integer, min(e.components, Nat(4)));
				input[inputCount].InputSlot = e.slot;
				input[inputCount].AlignedByteOffset = e.offset;
				input[inputCount].InputSlotClass = D3D10_INPUT_PER_VERTEX_DATA;
				input[inputCount].InstanceDataStepRate = 0;

				inputCount++;
			}
		}

		r = surface.dev()->CreateInputLayout(input, inputCount, codePtr, codeSize, &inputLayout.v);
		verify(shader->engine(), r, S("Failed to create an input layout: "));
	}


	D3DPixelShader::D3DPixelShader(PixelShader *input, D3DSurface &surface) {
		ComPtr<ID3D10Blob> code = compileShader(input, "ps_4_0");
		HRESULT r = surface.dev()->CreatePixelShader(code->GetBufferPointer(), code->GetBufferSize(), &shader.v);
		verify(input->engine(), r, S("Failed to create the pixel shader!"));
	}


	D3DPipeline::D3DPipeline(D3DVertexShader *vertex, D3DPixelShader *pixel, Pipeline *pipeline)
		: inputLayout(vertex->inputLayout),
		  vertex(vertex->shader), pixel(pixel->shader),
		  constantBuffers(pipeline->bufferCount(), ComPtr<ID3D10Buffer>()),
		  resourceViews(pipeline->textureCount(), ComPtr<ID3D10ShaderResourceView>()) {

		ComPtr<ID3D10Device> device;
		this->vertex->GetDevice(&device.v);

		for (size_t i = 0; i < pipeline->bufferCount(); i++) {
			const GcArray<Byte> *src = pipeline->buffer(i);

			D3D10_BUFFER_DESC desc;
			zeroMem(desc);
			desc.ByteWidth = UINT(src->count());
			desc.BindFlags = D3D10_BIND_CONSTANT_BUFFER;

			// #0 is special, it is for parameters:
			if (i == 0) {
				desc.Usage = D3D10_USAGE_DYNAMIC;
				desc.CPUAccessFlags = D3D10_CPU_ACCESS_WRITE;
			} else {
				desc.Usage = D3D10_USAGE_DEFAULT;
				desc.CPUAccessFlags = 0;
			}

			D3D10_SUBRESOURCE_DATA initData;
			zeroMem(initData);
			initData.pSysMem = src->v;

			HRESULT r = device->CreateBuffer(&desc, &initData, &constantBuffers[i].v);
			verify(pipeline->engine(), r, S("Failed to create a device buffer: "));
		}
	}

	void D3DPipeline::update(Graphics3D *graphics, Pipeline *pipeline, DirtyRange dirty) {
		// Note: We can (currently) assume that the buffer count never changes during the lifetime
		// of the pipeline object.

		ComPtr<ID3D10Device> device;
		this->vertex->GetDevice(&device.v);

		dirty.end = min(dirty.end, Nat(constantBuffers.size()));

		for (Nat i = dirty.begin; i < dirty.end; i++) {
			const GcArray<Byte> *src = pipeline->buffer(i);
			ComPtr<ID3D10Buffer> &buffer = constantBuffers[i];

			if (i == 0) {
				// #0 is special, use map API for it.
				void *bufData = null;
				HRESULT r = buffer->Map(D3D10_MAP_WRITE_DISCARD, 0, &bufData);
				verify(pipeline->engine(), r, S("Failed to map the buffer: "));
				memcpy(bufData, src->v, src->count());
				buffer->Unmap();
			} else {
				// Rest are DEFAULT, so use the UpdateSubresource method. Note that NULL updates the
				// entire thing. It seems that we can not update only a part of the subresource,
				// since the call does nothing if we pass anything other than NULL here.
				device->UpdateSubresource(buffer.v, 0, NULL, src->v, 0, 0);
			}
		}

		// Update textures. We simply check all of them.
		for (size_t i = 0; i < pipeline->textureCount(); i++) {
			Texture *texture = pipeline->texture(i);
			D3DTexture *d3d = (D3DTexture *)texture->forGraphics(graphics);

			resourceViews[i] = d3d->view;
		}
	}

}

#endif
