// TestApp1.cpp : Defines the entry point for the application.
//

#ifdef _WIN32
#include "Win32/stdafx.h"
#endif
#include <DK.h>

DKString ShaderStageNames(uint32_t s)
{
	DKArray<const char*> stages;
	if (s & (uint32_t)DKShaderStage::Vertex)
		stages.Add("Vertex");
	if (s & (uint32_t)DKShaderStage::TessellationControl)
		stages.Add("TessCtrl");
	if (s & (uint32_t)DKShaderStage::TessellationEvaluation)
		stages.Add("TessEval");
	if (s & (uint32_t)DKShaderStage::Geometry)
		stages.Add("Geometry");
	if (s & (uint32_t)DKShaderStage::Fragment)
		stages.Add("Fragment");
	if (s & (uint32_t)DKShaderStage::Compute)
		stages.Add("Compute");

	if (stages.IsEmpty())
		return "";

	DKString str = stages.Value(0);
	for (int i = 1; i < stages.Count(); ++i)
		str += DKString::Format(", %ls", stages.Value(i));
	return str;
}

const char* ShaderDataTypeStr(DKShaderDataType t)
{
	switch (t) {
	case DKShaderDataType::Unknown:				return "Unknown";
	case DKShaderDataType::None:				return "None";

	case DKShaderDataType::Struct:				return "Struct";
	case DKShaderDataType::Texture:				return "Texture";
	case DKShaderDataType::Sampler:				return "Sampler";

	case DKShaderDataType::Float:				return "Float";
	case DKShaderDataType::Float2:				return "Float2";
	case DKShaderDataType::Float3:				return "Float3";
	case DKShaderDataType::Float4:				return "Float4";

	case DKShaderDataType::Float2x2:			return "Float2x2";
	case DKShaderDataType::Float2x3:			return "Float2x3";
	case DKShaderDataType::Float2x4:			return "Float2x4";

	case DKShaderDataType::Float3x2:			return "Float3x2";
	case DKShaderDataType::Float3x3:			return "Float3x3";
	case DKShaderDataType::Float3x4:			return "Float3x4";

	case DKShaderDataType::Float4x2:			return "Float4x2";
	case DKShaderDataType::Float4x3:			return "Float4x3";
	case DKShaderDataType::Float4x4:			return "Float4x4";

	case DKShaderDataType::Half:				return "Half";
	case DKShaderDataType::Half2:				return "Half2";
	case DKShaderDataType::Half3:				return "Half3";
	case DKShaderDataType::Half4:				return "Half4";

	case DKShaderDataType::Half2x2:				return "Half2x2";
	case DKShaderDataType::Half2x3:				return "Half2x3";
	case DKShaderDataType::Half2x4:				return "Half2x4";

	case DKShaderDataType::Half3x2:				return "Half3x2";
	case DKShaderDataType::Half3x3:				return "Half3x3";
	case DKShaderDataType::Half3x4:				return "Half3x4";

	case DKShaderDataType::Half4x2:				return "Half4x2";
	case DKShaderDataType::Half4x3:				return "Half4x3";
	case DKShaderDataType::Half4x4:				return "Half4x4";

	case DKShaderDataType::Int:					return "Int";
	case DKShaderDataType::Int2:				return "Int2";
	case DKShaderDataType::Int3:				return "Int3";
	case DKShaderDataType::Int4:				return "Int4";

	case DKShaderDataType::UInt:				return "UInt";
	case DKShaderDataType::UInt2:				return "UInt2";
	case DKShaderDataType::UInt3:				return "UInt3";
	case DKShaderDataType::UInt4:				return "UInt4";

	case DKShaderDataType::Short:				return "Short";
	case DKShaderDataType::Short2:				return "Short2";
	case DKShaderDataType::Short3:				return "Short3";
	case DKShaderDataType::Short4:				return "Short4";

	case DKShaderDataType::UShort:				return "UShort";
	case DKShaderDataType::UShort2:				return "UShort2";
	case DKShaderDataType::UShort3:				return "UShort3";
	case DKShaderDataType::UShort4:				return "UShort4";

	case DKShaderDataType::Char:				return "Char";
	case DKShaderDataType::Char2:				return "Char2";
	case DKShaderDataType::Char3:				return "Char3";
	case DKShaderDataType::Char4:				return "Char4";

	case DKShaderDataType::UChar:				return "UChar";
	case DKShaderDataType::UChar2:				return "UChar2";
	case DKShaderDataType::UChar3:				return "UChar3";
	case DKShaderDataType::UChar4:				return "UChar4";

	case DKShaderDataType::Bool:				return "Bool";
	case DKShaderDataType::Bool2:				return "Bool2";
	case DKShaderDataType::Bool3:				return "Bool3";
	case DKShaderDataType::Bool4:				return "Bool4";
	}
	return "Error";
}

DKVertexDescriptor VertexDescriptorForVertexAttributes(const DKArray<DKShaderAttribute>& attrs)
{
    uint32_t vertexSize = 0;
    DKVertexDescriptor descriptor = {};

    for (uint32_t i = 0; i < attrs.Count(); ++i)
    {         
        if (const DKShaderAttribute& attr = attrs.Value(i); attr.enabled)
        {
            DKVertexAttributeDescriptor desc = {};

            switch (attr.type)
            {
               case DKShaderDataType::Float:
               case DKShaderDataType::Half:
               case DKShaderDataType::Int:
               case DKShaderDataType::UInt:
               case DKShaderDataType::Short:
               case DKShaderDataType::UShort:
               case DKShaderDataType::Char:
               case DKShaderDataType::UChar:
               case DKShaderDataType::Bool:
                   desc.format = DKVertexFormat::Float;
                   desc.offset = vertexSize;
                   vertexSize += 4;
                   break;
               case DKShaderDataType::Float2:
               case DKShaderDataType::Half2:
               case DKShaderDataType::Int2:
               case DKShaderDataType::UInt2:
               case DKShaderDataType::Short2:
               case DKShaderDataType::UShort2:
               case DKShaderDataType::Char2:
               case DKShaderDataType::UChar2:
               case DKShaderDataType::Bool2:
                   desc.format = DKVertexFormat::Float2;
                   desc.offset = vertexSize;
                   vertexSize += 8;
                   break;
               case DKShaderDataType::Float3:
               case DKShaderDataType::Half3:
               case DKShaderDataType::Int3:
               case DKShaderDataType::UInt3:
               case DKShaderDataType::Short3:
               case DKShaderDataType::UShort3:
               case DKShaderDataType::Char3:
               case DKShaderDataType::UChar3:
               case DKShaderDataType::Bool3:
                   desc.format = DKVertexFormat::Float3;
                   desc.offset = vertexSize;
                   vertexSize += 12;
                   break;
               case DKShaderDataType::Float4:
               case DKShaderDataType::Half4:
               case DKShaderDataType::Int4:
               case DKShaderDataType::UInt4:
               case DKShaderDataType::Short4:
               case DKShaderDataType::UShort4:
               case DKShaderDataType::Char4:
               case DKShaderDataType::UChar4:
               case DKShaderDataType::Bool4:
                   desc.format = DKVertexFormat::Float4;
                   desc.offset = vertexSize;
                   vertexSize += 16;
                   break;
               default:
                   DKASSERT_DESC(0, "Unsupported type!");
                   break;
            }
            desc.location = attr.location;

            descriptor.attributes.Add(desc);
        }
    }

    descriptor.layouts = { { DKVertexStepRate::Vertex, vertexSize}, };

    return descriptor;
};

void PrintShaderResource(const DKShaderResource& res, DKLogCategory c = DKLogCategory::Info)
{
	struct MemberPrinter
	{
		const DKShaderResource& res;
		int indent;
        DKLogCategory c;
		void operator()(const DKShaderResourceStruct& str) const
		{
			DKString indentStr = "";
			for (int i = 0; i < indent; ++i)
			{
				indentStr += "    ";
			}
			for (const DKShaderResourceStructMember& mem : str.members)
			{
				if (mem.count > 1)
				{
					DKLog(c, " %ls+ %ls[%d] (%s, Offset: %d, Stride: %d)",
						   (const wchar_t*)indentStr,
						   (const wchar_t*)mem.name,
						   mem.count,
						   ShaderDataTypeStr(mem.dataType),
						   mem.offset, mem.stride);

				}
				else
				{
					DKLog(c, " %ls+ %ls (%s, Offset: %d)",
						   (const wchar_t*)indentStr,
						   (const wchar_t*)mem.name,
						   ShaderDataTypeStr(mem.dataType),
						   mem.offset);
				}

				auto* p = res.structTypeMemberMap.Find(mem.typeInfoKey);
				if (p)
				{
					DKLog(c, " %ls  Struct \"%ls\"",
						   (const wchar_t*)indentStr,
						   (const wchar_t*)mem.typeInfoKey);
					MemberPrinter{ res, indent + 1, c}.operator()(p->value);
				}
			}
		}
	};


	if (res.count > 1)
		DKLog(c, "ShaderResource: %ls[%d] (set=%d, binding=%d, stages=%ls)",
			 (const wchar_t*)res.name, res.count, res.set, res.binding,
			 (const wchar_t*)ShaderStageNames(res.stages));
	else
		DKLog(c, "ShaderResource: %ls (set=%d, binding=%d, stages=%ls)",
			 (const wchar_t*)res.name, res.set, res.binding,
			 (const wchar_t*)ShaderStageNames(res.stages));

	const char* type = "Unknown (ERROR)";
	switch (res.type)
	{
	case DKShaderResource::TypeBuffer: type = "Buffer"; break;
	case DKShaderResource::TypeTexture:	type = "Texture"; break;
	case DKShaderResource::TypeSampler:	type = "Sampler"; break;
    case DKShaderResource::TypeTextureSampler: type = "SampledTexture"; break;
	}
	const char* access = "Unknown (ERROR)";
	switch (res.access)
	{
		case DKShaderResource::AccessReadOnly:	access = "ReadOnly"; break;
		case DKShaderResource::AccessWriteOnly:	access = "WriteOnly"; break;
		case DKShaderResource::AccessReadWrite:	access = "ReadWrite"; break;
	}

    if (res.type == DKShaderResource::TypeBuffer)
    {
        DKLog(c, " Type:%s, Access:%s, Enabled:%d, Size:%d",
               type,
               access,
               int(res.enabled),
               res.typeInfo.buffer.size);
    }
    else
    {
        DKLog(c, " Type:%s, Access:%s, Enabled:%d",
               type,
               access,
               int(res.enabled));
    }
	if (res.typeInfoKey.Length() > 0)
		DKLog(c, " Struct \"%ls\"", (const wchar_t*)res.typeInfoKey);
	if (res.type == DKShaderResource::TypeBuffer)
	{
		auto p = res.structTypeMemberMap.Find(res.typeInfoKey);
		if (p)
			MemberPrinter{ res, 1 , c}.operator()(p->value);
	}
}


class TestApp1 : public DKApplication
{
	DKResourcePool resourcePool;
	DKString userConfigPath;

    DKObject<DKGraphicsDevice> device;

	void PrintShaderReflection(const DKShader* shader, DKLogCategory c = DKLogCategory::Warning)
	{
		DKString stage = ShaderStageNames((uint32_t)shader->Stage());
		if (stage.Length() == 0)
			stage = "Unknown";

		DKLog(c, "=========================================================");
		DKLog(c, "Shader<%ls.SPIR-V>.InputAttributes: %d",
			(const wchar_t*)stage, shader->InputAttributes().Count());
		for (int i = 0; i < shader->InputAttributes().Count(); ++i)
		{
			const DKShaderAttribute& attr = shader->InputAttributes().Value(i);
			DKLog(c, "  [in] ShaderAttribute[%d]: \"%ls\" (location:%u)",
				  i, (const wchar_t*)attr.name, attr.location);
		}
		DKLog(c, "---------------------------------------------------------");
		DKLog(c, "Shader<%ls.SPIR-V>.OutputAttributes: %d",
			(const wchar_t*)stage, shader->OutputAttributes().Count());
		for (int i = 0; i < shader->OutputAttributes().Count(); ++i)
		{
			const DKShaderAttribute& attr = shader->OutputAttributes().Value(i);
			DKLog(c, "  [out] ShaderAttribute[%d]: \"%ls\" (location:%u)",
				  i, (const wchar_t*)attr.name, attr.location);
		}
		DKLog(c, "---------------------------------------------------------");
		DKLog(c, "Shader<%ls.SPIR-V>.Resources: %d",
			  (const wchar_t*)stage, shader->Resources().Count());
		for (auto& arg : shader->Resources())
			PrintShaderResource(arg, c);
		for (int i = 0; i < shader->PushConstantBufferLayouts().Count(); ++i)
		{
			const DKShaderPushConstantLayout& layout = shader->PushConstantBufferLayouts().Value(i);
			DKLog(c, " PushConstant:%d \"%ls\" (offset:%u, size:%u, stages:%ls)",
				  i, (const wchar_t*)layout.name, layout.offset, layout.size,
				  (const wchar_t*)ShaderStageNames(layout.stages));
		}
		DKLog(c, "=========================================================");
	}

	void PrintPipelineReflection(const DKPipelineReflection* reflection, DKLogCategory c = DKLogCategory::Error)
	{
		DKLog(c, "=========================================================");
		DKLog(c, "PipelineReflection.InputAttributes: %d", reflection->inputAttributes.Count());
		for (int i = 0; i < reflection->inputAttributes.Count(); ++i)
		{
			const DKShaderAttribute& attr = reflection->inputAttributes.Value(i);
			DKLog(c, "  [in] ShaderAttribute[%d]: \"%ls\" (location:%u)",
				  i, (const wchar_t*)attr.name, attr.location);
		}
		DKLog(c, "---------------------------------------------------------");
		DKLog(c, "PipelineReflection.Resources: %d", reflection->resources.Count());
		for (auto& arg : reflection->resources)
			PrintShaderResource(arg, c);
		for (int i = 0; i < reflection->pushConstantLayouts.Count(); ++i)
		{
			const DKShaderPushConstantLayout& layout = reflection->pushConstantLayouts.Value(i);
			DKLog(c, " PushConstant:%d \"%ls\" (offset:%u, size:%u, stages:%ls)",
				  i, (const wchar_t*)layout.name, layout.offset, layout.size,
				  (const wchar_t*)ShaderStageNames(layout.stages));
		}
		DKLog(c, "=========================================================");
	}

    void TestRenderPipelineReflection(const DKString& path)
    {
        DKObject<DKData> vertData = resourcePool.LoadResourceData(path + ".vert.spv");
        DKObject<DKData> fragData = resourcePool.LoadResourceData(path + ".frag.spv");

        DKASSERT(vertData);
        DKASSERT(fragData);

        DKShader vertShader(vertData);
        DKShader fragShader(fragData);

		PrintShaderReflection(&vertShader, DKLogCategory::Warning);
		PrintShaderReflection(&fragShader, DKLogCategory::Warning);

        DKObject<DKShaderModule> vertShaderModule = device->CreateShaderModule(&vertShader);
        DKObject<DKShaderModule> fragShaderModule = device->CreateShaderModule(&fragShader);

        DKObject<DKShaderFunction> vertShaderFunction = vertShaderModule->CreateFunction(vertShaderModule->FunctionNames().Value(0));
        DKObject<DKShaderFunction> fragShaderFunction = fragShaderModule->CreateFunction(fragShaderModule->FunctionNames().Value(0));

        // setup dummy vertex-descriptor
        DKVertexDescriptor vertexDescriptor = VertexDescriptorForVertexAttributes(vertShaderFunction->StageInputAttributes());

        // setup rendering pipeline state object (PSO)
        DKRenderPipelineDescriptor pipelineDescriptor = {};
        pipelineDescriptor.vertexFunction = vertShaderFunction;
        pipelineDescriptor.fragmentFunction = fragShaderFunction;
        pipelineDescriptor.colorAttachments.Resize(1);
        pipelineDescriptor.colorAttachments.Value(0).pixelFormat = DKPixelFormat::RGBA8Unorm;
        pipelineDescriptor.depthStencilAttachmentPixelFormat = DKPixelFormat::Invalid; // no depth buffer
        pipelineDescriptor.vertexDescriptor = vertexDescriptor;
        pipelineDescriptor.primitiveTopology = DKPrimitiveType::Triangle;
        pipelineDescriptor.frontFace = DKFrontFace::CCW;
        pipelineDescriptor.triangleFillMode = DKTriangleFillMode::Fill;
        pipelineDescriptor.depthClipMode = DKDepthClipMode::Clip;
        pipelineDescriptor.cullMode = DKCullMode::None;
        pipelineDescriptor.rasterizationEnabled = true;

        DKPipelineReflection reflection;
        DKObject<DKRenderPipelineState> pipelineState = device->CreateRenderPipeline(pipelineDescriptor, &reflection);
        if (pipelineState)
        {
			PrintPipelineReflection(&reflection, DKLogCategory::Error);
        }
        else
        {
            DKLogE("CreateRenderPipeline failed.");
        }
    }

    void TestComputePipelineReflection(const DKString& path)
    {
        DKObject<DKData> shaderData = resourcePool.LoadResourceData(path + ".comp.spv");

        DKASSERT(shaderData);

        DKShader shader(shaderData);

		PrintShaderReflection(&shader, DKLogCategory::Warning);

        DKObject<DKShaderModule> shaderModule = device->CreateShaderModule(&shader);
        DKObject<DKShaderFunction> shaderFunction = shaderModule->CreateFunction(shaderModule->FunctionNames().Value(0));


        // setup pipeline state object (PSO)
        DKComputePipelineDescriptor pipelineDescriptor = {};
        pipelineDescriptor.computeFunction = shaderFunction;

        DKPipelineReflection reflection;
        DKObject<DKComputePipelineState> pipelineState = device->CreateComputePipeline(pipelineDescriptor, &reflection);
        if (pipelineState)
        {
			PrintPipelineReflection(&reflection, DKLogCategory::Error);
        }
        else
        {
            DKLogE("CreateComputePipeline failed.");
        }
    }

public:
	void OnInitialize(void) override
	{
		DKLogD("%s", DKGL_FUNCTION_NAME);
		DKLogD("NumCores: %d", DKNumberOfCpuCores());
		DKLogD("NumProcessors: %d", DKNumberOfProcessors());

		DKString resPath = DefaultPath(SystemPath::AppResource);
		resPath = resPath.FilePathStringByAppendingPath("Data");
		DKLog("resPath: %ls", (const wchar_t*)resPath);
		resourcePool.AddLocatorForPath(resPath);

		userConfigPath = DefaultPath(SystemPath::AppExecutable).FilePathStringByAppendingPath("app.config.xml");
		int numItemsImported = DKPropertySet::DefaultSet().Import(userConfigPath, true);
		if (numItemsImported >= 0)
			DKLogI("UserSettings: %ls (%d items imported)", (const wchar_t*)userConfigPath, numItemsImported);

        device = DKGraphicsDevice::SharedInstance();

        TestRenderPipelineReflection("shaders/conditionalrender/model");
        TestComputePipelineReflection("shaders/computeshader/emboss");

		Terminate(0);
	}
	void OnTerminate(void) override
	{
		DKLogD("%s", DKGL_FUNCTION_NAME);

        device = NULL;

		int numItemsExported = DKPropertySet::DefaultSet().Export(userConfigPath, true);
		DKLogI("Setting saved: %ls (%d items exported)", (const wchar_t*)userConfigPath, numItemsExported);

		DKLogI("Memory Pool Statistics");
		size_t numBuckets = DKMemoryPoolNumberOfBuckets();
		DKMemoryPoolBucketStatus* buckets = new DKMemoryPoolBucketStatus[numBuckets];
		DKMemoryPoolQueryAllocationStatus(buckets, numBuckets);
		size_t usedBytes = 0;
		for (int i = 0; i < numBuckets; ++i)
		{
			if (buckets[i].totalChunks > 0)
			{
				DKLogI("--> %5lu:  %5lu/%5lu, usage: %.1f%%, used: %.1fKB, total: %.1fKB",
					   buckets[i].chunkSize,
					   buckets[i].usedChunks, buckets[i].totalChunks,
					   double(buckets[i].usedChunks) / double(buckets[i].totalChunks) * 100.0,
					   double(buckets[i].chunkSize * buckets[i].usedChunks) / 1024.0,
					   double(buckets[i].chunkSize * buckets[i].totalChunks) / 1024.0
				);
				usedBytes += buckets[i].chunkSize * buckets[i].usedChunks;
			}
		}
		DKLogI("MemoryPool Usage: %.1fMB / %.1fMB", double(usedBytes) / (1024 * 1024), double(DKMemoryPoolSize()) / (1024 * 1024));
		delete[] buckets;
	}
};


#ifdef _WIN32
int APIENTRY wWinMain(_In_ HINSTANCE hInstance,
					  _In_opt_ HINSTANCE hPrevInstance,
					  _In_ LPWSTR    lpCmdLine,
					  _In_ int       nCmdShow)
#else
int main(int argc, const char * argv[])
#endif
{
	TestApp1 app;
	DKPropertySet::SystemConfig().SetValue("AppDelegate", "AppDelegate");
	DKPropertySet::SystemConfig().SetValue("GraphicsAPI", "Vulkan");
	return app.Run();
}