#include "Global_Header.h"
#include "Vulkan_Header.h"
#include "Windows_Misc.h"
#include <math.h>

extern VulkanSystem _VkSystem;

static void BeginRenderPass(VkRenderPass renderPass, VkFramebuffer frameBuffer, VkCommandBuffer commandBuffer,
	float r, float g, float b)
{
	VkClearValue clearValue = {};
	clearValue.color = { r,g,b,0 };

	VkRenderPassBeginInfo passInfo = {};
	passInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
	passInfo.renderPass = renderPass;
	passInfo.framebuffer = frameBuffer;
	passInfo.pClearValues = &clearValue;
	passInfo.clearValueCount = 1;
	passInfo.renderArea.extent.width = SCREEN_WIDTH;
	passInfo.renderArea.extent.height = SCREEN_HEIGHT;

	vkCmdBeginRenderPass(commandBuffer, &passInfo, VK_SUBPASS_CONTENTS_BEGIN_RANGE);
}

static int SubmitToQueue(VkQueue queue, VkCommandBuffer commandBuffer, VkFence fence, VkSemaphore waitSemap, VkSemaphore signalSemap)
{
	VkSubmitInfo submitInfo = {};
	submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
	submitInfo.commandBufferCount = 1;
	submitInfo.pCommandBuffers = &commandBuffer;
	submitInfo.pWaitSemaphores = &waitSemap;
	submitInfo.waitSemaphoreCount = 1;
	submitInfo.pSignalSemaphores = &signalSemap;
	submitInfo.signalSemaphoreCount = 1;

	VkPipelineStageFlags pipeFlags = VK_PIPELINE_STAGE_TRANSFER_BIT;
	submitInfo.pWaitDstStageMask = &pipeFlags;

	VkResult res = vkQueueSubmit(queue, 1, &submitInfo, fence);

	return res;
}

static int FindMemoryType(uint32_t memoryTypeBits, uint32_t properties)
{
	// Get Physical Device Properties
	VkPhysicalDeviceMemoryProperties devMemProp = {};
	vkGetPhysicalDeviceMemoryProperties(_VkSystem.physicalDevice, &devMemProp);

	// Check which index of the memory properties supports HOST_VISIBLE?
	int typeIndex = -1;
	for (uint32_t i = 0; i < devMemProp.memoryTypeCount; i++)
	{
		if ((memoryTypeBits & (1 << i)) && ((devMemProp.memoryTypes[i].propertyFlags & properties) == properties))
		{
			typeIndex = i;
			break;
		}
	}

	return typeIndex;
}

// Not too horrible
static int SetupBuffers()
{
	// Vertex Data
	float vertData[] = {
		-1.0f,1.0f , -1.0f,-1.0f, 1.0f,-1.0f,
		-1.0f,1.0f , 1.0f,-1.0f, 1.0f,1.0f
	};

	// Setup The Vertex Data Buffer information
	VkBufferCreateInfo bufferInfo = {};
	bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
	bufferInfo.usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT;
	bufferInfo.size = sizeof(vertData);
	bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

	// Create Vertex Data Buffer
	VkBuffer buffer;
	VkResult res = vkCreateBuffer(_VkSystem.device, &bufferInfo, 0, &buffer);
	if (res != VK_SUCCESS)
		return 1;

	// Get the memory requirements for the buffer
	VkMemoryRequirements memoryReq;
	vkGetBufferMemoryRequirements(_VkSystem.device, buffer, &memoryReq);

	uint32_t typeIndex = FindMemoryType(memoryReq.memoryTypeBits, (VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT));
	if (typeIndex == -1)
		DebugBreak();

	// Setup Memory Allocation Information
	VkMemoryAllocateInfo allocInfo = {};
	allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
	allocInfo.allocationSize = memoryReq.size;
	allocInfo.memoryTypeIndex = typeIndex;

	// Allocate device memory
	VkDeviceMemory devMemory = {};
	res = vkAllocateMemory(_VkSystem.device, &allocInfo, 0, &devMemory);
	if (res != VK_SUCCESS)
		return 1;

	// Bind the Vertex Data Buffer to our Device Memory
	res = vkBindBufferMemory(_VkSystem.device, buffer, devMemory, 0);
	if (res != VK_SUCCESS)
		return 1;


	// Map the Device Memory so we can write to it
	void* mappedMemory = 0;
	res = vkMapMemory(_VkSystem.device, devMemory, 0, bufferInfo.size, 0, &mappedMemory);
	if (res != VK_SUCCESS)
		return 1;

	// Copy to Device Memory
	memcpy(mappedMemory, vertData, bufferInfo.size);

	//// Setup Memory Range needed to Flush the data
	//VkMappedMemoryRange memoryRange = {};
	//memoryRange.sType = VK_STRUCTURE_TYPE_MAPPED_MEMORY_RANGE;
	//memoryRange.memory = devMemory;
	//memoryRange.size = VK_WHOLE_SIZE;

	//// Flush Mapped Memory
	//res = vkFlushMappedMemoryRanges(_VkSystem.device, 1, &memoryRange);
	//if (res != VK_SUCCESS)
	//	return 1;

	// Unmap Device Memory
	vkUnmapMemory(_VkSystem.device, devMemory);

	_VkSystem.buffer = buffer;

	return res;
}

// Create a Shader Module from a SPIR-V Shader
static int CreateShaderModule(const wchar_t* file, VkShaderModule* pShaderModule)
{
	// Load Shader (SPIR-V) to Memory temporarily
	unsigned long size;
	void* shaderFile = Win32_FileLoader(file, &size);

	VkShaderModuleCreateInfo moduleInfo = {};
	moduleInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
	moduleInfo.codeSize = size;
	moduleInfo.pCode = (uint32_t*)shaderFile;

	// Create Shader Module
	VkShaderModule shaderModule;
	VkResult res = vkCreateShaderModule(_VkSystem.device, &moduleInfo, 0, &shaderModule);
	if (res != VK_SUCCESS)
		return 1;

	// Free Memory used for the Shader (SPIR-V)
	Win32_Free(shaderFile);

	*pShaderModule = shaderModule;

	return res;
}

// Setup Shader Stage Info (For convinience)
static void SetupShaderStageInfo(VkPipelineShaderStageCreateInfo* stageInfo, VkShaderModule shaderModule, const char* entryPoint, VkShaderStageFlagBits flagBits)
{
	stageInfo->sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
	stageInfo->stage = flagBits;
	stageInfo->pName = "main";
	stageInfo->module = shaderModule;
}

// Horrible
static int SetupGraphicsPipeline()
{
	VkResult res;

	VkShaderModule vertexModule;
	if (CreateShaderModule(L"vert.spv", &vertexModule) != VK_SUCCESS)
		return 1;

	// Create Fragment Shader Module
	VkShaderModule fragModule;
	if (CreateShaderModule(L"frag.spv", &fragModule) != VK_SUCCESS)
		return 1;

	// Create 2 Shader Stage Create Info for Vertex and Fragment shaders
	VkPipelineShaderStageCreateInfo stageInfo[2] = {};

	// Setup Stage info for the Vertex Shader
	SetupShaderStageInfo(&stageInfo[0], vertexModule, "main", VK_SHADER_STAGE_VERTEX_BIT);

	// Setup Stage info for the Fragment Shader
	SetupShaderStageInfo(&stageInfo[1], fragModule, "main", VK_SHADER_STAGE_FRAGMENT_BIT);

	VkPipelineLayoutCreateInfo layoutInfo = {};
	layoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;


	// Create the Pipeline Layout
	VkPipelineLayout pipeLayout;
	res = vkCreatePipelineLayout(_VkSystem.device, &layoutInfo, 0, &pipeLayout);
	if (res != VK_SUCCESS)
		return 1;

	// Viewport and Scissors should be the Window Width/Height
	VkViewport viewPort = {};
	viewPort.width = SCREEN_WIDTH;
	viewPort.height = SCREEN_HEIGHT;

	VkRect2D scissors = {};
	scissors.extent.width = SCREEN_WIDTH;
	scissors.extent.height = SCREEN_HEIGHT;

	// Setup Viewport state information
	VkPipelineViewportStateCreateInfo viewportInfo = {};
	viewportInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
	viewportInfo.viewportCount = viewportInfo.scissorCount = 1;
	viewportInfo.pViewports = (const VkViewport*)&viewportInfo;
	viewportInfo.pScissors = &scissors;

	// Setup Rasterization State Information
	VkPipelineRasterizationStateCreateInfo rasterInfo = {};
	rasterInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
	rasterInfo.polygonMode = VK_POLYGON_MODE_FILL;
	rasterInfo.cullMode = VK_CULL_MODE_NONE;
	rasterInfo.frontFace = VK_FRONT_FACE_CLOCKWISE;
	rasterInfo.lineWidth = 1.0f;

	// Setup Input Assembly State Information
	VkPipelineInputAssemblyStateCreateInfo assemblyInfo = {};
	assemblyInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
	assemblyInfo.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

	// Setup Vertex Input Attribute Description (2 32bit floats at location=0)
	VkVertexInputAttributeDescription vertAttribDesc = {};
	vertAttribDesc.binding = 0;
	vertAttribDesc.location = 0;
	vertAttribDesc.offset = 0;
	vertAttribDesc.format = VK_FORMAT_R32G32_SFLOAT;

	// Setup Vertex Input Binding Description
	VkVertexInputBindingDescription vertBindDesc = {};
	vertBindDesc.binding = 0;
	vertBindDesc.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
	vertBindDesc.stride = sizeof(float) * 2;

	// Setup Vertex Input State Information
	VkPipelineVertexInputStateCreateInfo vertInfo = {};
	vertInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
	vertInfo.pVertexAttributeDescriptions = &vertAttribDesc;
	vertInfo.vertexAttributeDescriptionCount = 1;
	vertInfo.pVertexBindingDescriptions = &vertBindDesc;
	vertInfo.vertexBindingDescriptionCount = 1;

	VkPipelineMultisampleStateCreateInfo multiInfo = {};
	multiInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
	multiInfo.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

	VkPipelineColorBlendAttachmentState blendAttachment = {};

	VkPipelineColorBlendStateCreateInfo blendInfo = {};
	blendInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
	blendInfo.attachmentCount = 1;
	blendInfo.pAttachments = &blendAttachment;

	// Setup Graphics Pipeline Information
	VkGraphicsPipelineCreateInfo pipelineInfo = {};
	pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
	pipelineInfo.renderPass = _VkSystem.renderPass;
	pipelineInfo.stageCount = 2;
	pipelineInfo.pStages = stageInfo;
	pipelineInfo.pVertexInputState = &vertInfo;
	pipelineInfo.pInputAssemblyState = &assemblyInfo;
	pipelineInfo.pRasterizationState = &rasterInfo;
	pipelineInfo.layout = pipeLayout;
	pipelineInfo.pViewportState = &viewportInfo;
	pipelineInfo.pColorBlendState = &blendInfo;
	pipelineInfo.pMultisampleState = &multiInfo;

	// Create Graphics Pipeline
	VkPipeline graphicPipeline;
	res = vkCreateGraphicsPipelines(_VkSystem.device, 0, 1, &pipelineInfo, 0, &graphicPipeline);
	if (res != VK_SUCCESS)
		return 1;

	_VkSystem.graphicsPipeline = graphicPipeline;

	return res;
}

int QueuePresent(VkQueue queue, VkSwapchainKHR swapchain, VkSemaphore semaphoreSubmit, uint32_t imageIndex)
{
	VkPresentInfoKHR presentInfo = {};
	presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
	presentInfo.pSwapchains = &swapchain;
	presentInfo.swapchainCount = 1;
	presentInfo.pWaitSemaphores = &semaphoreSubmit;
	presentInfo.waitSemaphoreCount = 1;
	presentInfo.pImageIndices = &_VkSystem.imageIndex;

	VkResult res = vkQueuePresentKHR(queue, &presentInfo);

	return res;
}

int Vulkan_Render(float color)
{
	VkResult res;

	// Wait for Fence
	vkWaitForFences(_VkSystem.device, 1, &_VkSystem.fence, 1, UINT64_MAX);

	// Acquire the next image to work at (this is technically allowed to fail, but we need to restart the loop then)
	res = vkAcquireNextImageKHR(_VkSystem.device, _VkSystem.swapchainKHR, UINT64_MAX, _VkSystem.semaphoreNextImage, VK_NULL_HANDLE, &_VkSystem.imageIndex);
	if (res != VK_SUCCESS)
		return 1;

	// Reset Fence
	vkResetFences(_VkSystem.device, 1, &_VkSystem.fence);

	VkCommandBufferBeginInfo beginInfo = {};
	beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
	beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

	// Begin CommandBuffer Recording
	vkBeginCommandBuffer(_VkSystem.commandBuffer, &beginInfo);

	static int temp;
	if (!temp)
	{
		SetupGraphicsPipeline();
		SetupBuffers();
		temp = 1;
	}

	VkDeviceSize devSize = 0;

	// Bind Graphics Pipeline
	vkCmdBindPipeline(_VkSystem.commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, _VkSystem.graphicsPipeline);
	// Bind Vertex Buffers
	vkCmdBindVertexBuffers(_VkSystem.commandBuffer, 0, 1, &_VkSystem.buffer, &devSize);

	// Begin RenderPass
	BeginRenderPass(_VkSystem.renderPass, _VkSystem.frameBuffers[_VkSystem.imageIndex].frameBuffer,
		_VkSystem.commandBuffer, 0.5f, 0.1f, 0.8f);

	// Draw Vertex Buffer
	vkCmdDraw(_VkSystem.commandBuffer, 6, 1, 0, 0);

	// End RenderPass
	vkCmdEndRenderPass(_VkSystem.commandBuffer);

	// End CommandBuffer Recording
	res = vkEndCommandBuffer(_VkSystem.commandBuffer);
	if (res != VK_SUCCESS)
		return 1;

	if (SubmitToQueue(_VkSystem.queue.queue, _VkSystem.commandBuffer, _VkSystem.fence,
		_VkSystem.semaphoreNextImage, _VkSystem.queue.semaphoreSubmit) != VK_SUCCESS)
		return 1;

	QueuePresent(_VkSystem.queue.queue, _VkSystem.swapchainKHR, _VkSystem.queue.semaphoreSubmit, _VkSystem.imageIndex);
}