#include "graphics-pipeline_vulkan.hpp"

#include <fstream>
#include <stdexcept>
#include <iostream>

#include "vulkan-utils.hpp"

using namespace std;

// TODO: Remove any instances of cout and instead throw exceptions

GraphicsPipeline_Vulkan::GraphicsPipeline_Vulkan(VkPhysicalDevice physicalDevice, VkDevice device,
      VkRenderPass renderPass, Viewport viewport, int vertexSize) {
   this->physicalDevice = physicalDevice;
   this->device = device;
   this->renderPass = renderPass;
   this->viewport = viewport;

   // Since there is only one array of vertex data, we use binding = 0
   // I'll probably do that for the foreseeable future
   // I can calculate the stride myself given info about all the varying attributes
   this->bindingDescription.binding = 0;
   this->bindingDescription.stride = vertexSize;
   this->bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
}

GraphicsPipeline_Vulkan::~GraphicsPipeline_Vulkan() {
}

void GraphicsPipeline_Vulkan::createVertexBuffer(const void* bufferData, int vertexSize,
      VkCommandPool commandPool, VkQueue graphicsQueue) {
   VkDeviceSize bufferSize = numVertices * vertexSize;
   VkDeviceSize bufferCapacity = vertexCapacity * vertexSize;

   VkBuffer stagingBuffer;
   VkDeviceMemory stagingBufferMemory;
   VulkanUtils::createBuffer(device, physicalDevice, bufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
      stagingBuffer, stagingBufferMemory);

   void* data;
   vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);
   memcpy(data, bufferData, (size_t) bufferSize);
   vkUnmapMemory(device, stagingBufferMemory);

   VulkanUtils::createBuffer(device, physicalDevice, bufferCapacity,
      VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
      VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, vertexBuffer, vertexBufferMemory);

   VulkanUtils::copyBuffer(device, commandPool, stagingBuffer, vertexBuffer, 0, 0, bufferSize, graphicsQueue);

   vkDestroyBuffer(device, stagingBuffer, nullptr);
   vkFreeMemory(device, stagingBufferMemory, nullptr);
}

void GraphicsPipeline_Vulkan::createIndexBuffer(const void* bufferData, int indexSize,
      VkCommandPool commandPool, VkQueue graphicsQueue) {
   VkDeviceSize bufferSize = numIndices * indexSize;
   VkDeviceSize bufferCapacity = indexCapacity * indexSize;

   VkBuffer stagingBuffer;
   VkDeviceMemory stagingBufferMemory;
   VulkanUtils::createBuffer(device, physicalDevice, bufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
      stagingBuffer, stagingBufferMemory);

   void* data;
   vkMapMemory(device, stagingBufferMemory, 0, bufferSize, 0, &data);
   memcpy(data, bufferData, (size_t) bufferSize);
   vkUnmapMemory(device, stagingBufferMemory);

   VulkanUtils::createBuffer(device, physicalDevice, bufferCapacity,
      VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
      VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, indexBuffer, indexBufferMemory);

   VulkanUtils::copyBuffer(device, commandPool, stagingBuffer, indexBuffer, 0, 0, bufferSize, graphicsQueue);

   vkDestroyBuffer(device, stagingBuffer, nullptr);
   vkFreeMemory(device, stagingBufferMemory, nullptr);
}

void GraphicsPipeline_Vulkan::addAttribute(VkFormat format, size_t offset) {
   VkVertexInputAttributeDescription attributeDesc = {};

   attributeDesc.binding = 0;
   attributeDesc.location = this->attributeDescriptions.size();
   attributeDesc.format = format;
   attributeDesc.offset = offset;

   this->attributeDescriptions.push_back(attributeDesc);
}

void GraphicsPipeline_Vulkan::addDescriptorInfo(VkDescriptorType type, VkShaderStageFlags stageFlags, vector<VkDescriptorBufferInfo>* bufferData) {
   this->descriptorInfoList.push_back({ type, stageFlags, bufferData, nullptr });
}

void GraphicsPipeline_Vulkan::addDescriptorInfo(VkDescriptorType type, VkShaderStageFlags stageFlags, VkDescriptorImageInfo* imageData) {
   this->descriptorInfoList.push_back({ type, stageFlags, nullptr, imageData });
}

void GraphicsPipeline_Vulkan::createPipeline(string vertShaderFile, string fragShaderFile) {
   vector<char> vertShaderCode = readFile(vertShaderFile);
   vector<char> fragShaderCode = readFile(fragShaderFile);

   VkShaderModule vertShaderModule = createShaderModule(vertShaderCode);
   VkShaderModule fragShaderModule = createShaderModule(fragShaderCode);

   VkPipelineShaderStageCreateInfo vertShaderStageInfo = {};
   vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
   vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
   vertShaderStageInfo.module = vertShaderModule;
   vertShaderStageInfo.pName = "main";

   VkPipelineShaderStageCreateInfo fragShaderStageInfo = {};
   fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
   fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
   fragShaderStageInfo.module = fragShaderModule;
   fragShaderStageInfo.pName = "main";

   VkPipelineShaderStageCreateInfo shaderStages[] = { vertShaderStageInfo, fragShaderStageInfo };

   VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};
   vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;

   vertexInputInfo.vertexBindingDescriptionCount = 1;
   vertexInputInfo.vertexAttributeDescriptionCount = static_cast<uint32_t>(this->attributeDescriptions.size());
   vertexInputInfo.pVertexBindingDescriptions = &this->bindingDescription;
   vertexInputInfo.pVertexAttributeDescriptions = this->attributeDescriptions.data();

   VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
   inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
   inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
   inputAssembly.primitiveRestartEnable = VK_FALSE;

   VkViewport viewport = {};
   viewport.x = (float)this->viewport.x;
   viewport.y = (float)this->viewport.y;
   viewport.width = (float)this->viewport.width;
   viewport.height = (float)this->viewport.height;
   viewport.minDepth = 0.0f;
   viewport.maxDepth = 1.0f;

   VkRect2D scissor = {};
   scissor.offset = { 0, 0 };
   scissor.extent = { (uint32_t)this->viewport.width, (uint32_t)this->viewport.height };

   VkPipelineViewportStateCreateInfo viewportState = {};
   viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
   viewportState.viewportCount = 1;
   viewportState.pViewports = &viewport;
   viewportState.scissorCount = 1;
   viewportState.pScissors = &scissor;

   VkPipelineRasterizationStateCreateInfo rasterizer = {};
   rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
   rasterizer.depthClampEnable = VK_FALSE;
   rasterizer.rasterizerDiscardEnable = VK_FALSE;
   rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
   rasterizer.lineWidth = 1.0f;
   rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
   rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
   rasterizer.depthBiasEnable = VK_FALSE;

   VkPipelineMultisampleStateCreateInfo multisampling = {};
   multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
   multisampling.sampleShadingEnable = VK_FALSE;
   multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

   VkPipelineColorBlendAttachmentState colorBlendAttachment = {};
   colorBlendAttachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
   colorBlendAttachment.blendEnable = VK_TRUE;
   colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
   colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
   colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
   colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;
   colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
   colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;

   VkPipelineColorBlendStateCreateInfo colorBlending = {};
   colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
   colorBlending.logicOpEnable = VK_FALSE;
   colorBlending.logicOp = VK_LOGIC_OP_COPY;
   colorBlending.attachmentCount = 1;
   colorBlending.pAttachments = &colorBlendAttachment;
   colorBlending.blendConstants[0] = 0.0f;
   colorBlending.blendConstants[1] = 0.0f;
   colorBlending.blendConstants[2] = 0.0f;
   colorBlending.blendConstants[3] = 0.0f;

   VkPipelineDepthStencilStateCreateInfo depthStencil = {};
   depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
   depthStencil.depthTestEnable = VK_TRUE;
   depthStencil.depthWriteEnable = VK_TRUE;
   depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
   depthStencil.depthBoundsTestEnable = VK_FALSE;
   depthStencil.minDepthBounds = 0.0f;
   depthStencil.maxDepthBounds = 1.0f;
   depthStencil.stencilTestEnable = VK_FALSE;
   depthStencil.front = {};
   depthStencil.back = {};

   VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};
   pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
   pipelineLayoutInfo.setLayoutCount = 1;
   pipelineLayoutInfo.pSetLayouts = &this->descriptorSetLayout;
   pipelineLayoutInfo.pushConstantRangeCount = 0;

   if (vkCreatePipelineLayout(this->device, &pipelineLayoutInfo, nullptr, &this->pipelineLayout) != VK_SUCCESS) {
      throw runtime_error("failed to create pipeline layout!");
   }

   VkGraphicsPipelineCreateInfo pipelineInfo = {};
   pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
   pipelineInfo.stageCount = 2;
   pipelineInfo.pStages = shaderStages;
   pipelineInfo.pVertexInputState = &vertexInputInfo;
   pipelineInfo.pInputAssemblyState = &inputAssembly;
   pipelineInfo.pViewportState = &viewportState;
   pipelineInfo.pRasterizationState = &rasterizer;
   pipelineInfo.pMultisampleState = &multisampling;
   pipelineInfo.pDepthStencilState = &depthStencil;
   pipelineInfo.pColorBlendState = &colorBlending;
   pipelineInfo.pDynamicState = nullptr;
   pipelineInfo.layout = this->pipelineLayout;
   pipelineInfo.renderPass = this->renderPass;
   pipelineInfo.subpass = 0;
   pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
   pipelineInfo.basePipelineIndex = -1;

   if (vkCreateGraphicsPipelines(this->device, VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &this->pipeline) != VK_SUCCESS) {
      throw runtime_error("failed to create graphics pipeline!");
   }

   vkDestroyShaderModule(this->device, vertShaderModule, nullptr);
   vkDestroyShaderModule(this->device, fragShaderModule, nullptr);
}

void GraphicsPipeline_Vulkan::createDescriptorSetLayout() {
   vector<VkDescriptorSetLayoutBinding> bindings(this->descriptorInfoList.size());

   for (size_t i = 0; i < bindings.size(); i++) {
      bindings[i].binding = i;
      bindings[i].descriptorCount = 1;
      bindings[i].descriptorType = this->descriptorInfoList[i].type;
      bindings[i].stageFlags = this->descriptorInfoList[i].stageFlags;
      bindings[i].pImmutableSamplers = nullptr;
   }

   VkDescriptorSetLayoutCreateInfo layoutInfo = {};
   layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
   layoutInfo.bindingCount = static_cast<uint32_t>(bindings.size());
   layoutInfo.pBindings = bindings.data();

   if (vkCreateDescriptorSetLayout(this->device, &layoutInfo, nullptr, &this->descriptorSetLayout) != VK_SUCCESS) {
      throw runtime_error("failed to create descriptor set layout!");
   }
}

void GraphicsPipeline_Vulkan::createDescriptorPool(vector<VkImage>& swapChainImages) {
   vector<VkDescriptorPoolSize> poolSizes(this->descriptorInfoList.size());

   for (size_t i = 0; i < poolSizes.size(); i++) {
      poolSizes[i].type = this->descriptorInfoList[i].type;
      poolSizes[i].descriptorCount = static_cast<uint32_t>(swapChainImages.size());
   }

   VkDescriptorPoolCreateInfo poolInfo = {};
   poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
   poolInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
   poolInfo.pPoolSizes = poolSizes.data();
   poolInfo.maxSets = static_cast<uint32_t>(swapChainImages.size());

   if (vkCreateDescriptorPool(this->device, &poolInfo, nullptr, &this->descriptorPool) != VK_SUCCESS) {
      throw runtime_error("failed to create descriptor pool!");
   }
}

void GraphicsPipeline_Vulkan::createDescriptorSets(vector<VkImage>& swapChainImages) {
   vector<VkDescriptorSetLayout> layouts(swapChainImages.size(), this->descriptorSetLayout);

   VkDescriptorSetAllocateInfo allocInfo = {};
   allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
   allocInfo.descriptorPool = this->descriptorPool;
   allocInfo.descriptorSetCount = static_cast<uint32_t>(swapChainImages.size());
   allocInfo.pSetLayouts = layouts.data();

   this->descriptorSets.resize(swapChainImages.size());
   if (vkAllocateDescriptorSets(device, &allocInfo, this->descriptorSets.data()) != VK_SUCCESS) {
      throw runtime_error("failed to allocate descriptor sets!");
   }

   for (size_t i = 0; i < swapChainImages.size(); i++) {
      vector<VkWriteDescriptorSet> descriptorWrites(this->descriptorInfoList.size());

      for (size_t j = 0; j < descriptorWrites.size(); j++) {
         descriptorWrites[j].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
         descriptorWrites[j].dstSet = this->descriptorSets[i];
         descriptorWrites[j].dstBinding = j;
         descriptorWrites[j].dstArrayElement = 0;
         descriptorWrites[j].descriptorType = this->descriptorInfoList[j].type;
         descriptorWrites[j].descriptorCount = 1;
         descriptorWrites[j].pBufferInfo = nullptr;
         descriptorWrites[j].pImageInfo = nullptr;
         descriptorWrites[j].pTexelBufferView = nullptr;

         switch (descriptorWrites[j].descriptorType) {
            case VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER:
               descriptorWrites[j].pBufferInfo = &(*this->descriptorInfoList[j].bufferDataList)[i];
               break;
            case VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER:
               descriptorWrites[j].pImageInfo = this->descriptorInfoList[j].imageData;
               break;
            default:
               cout << "Unknown descriptor type: " << descriptorWrites[j].descriptorType << endl;
         }
      }

      vkUpdateDescriptorSets(this->device, static_cast<uint32_t>(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);
   }
}

void GraphicsPipeline_Vulkan::createRenderCommands(VkCommandBuffer& commandBuffer, uint32_t currentImage) {
   vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline);
   vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout, 0, 1,
      &descriptorSets[currentImage], 0, nullptr);

   // TODO: Implement once I add vertex and index buffers to the pipeline
   /*
   VkBuffer vertexBuffers[] = { info.vertexBuffer };
   VkDeviceSize offsets[] = { 0 };
   vkCmdBindVertexBuffers(commandBuffers[currentImage], 0, 1, vertexBuffers, offsets);

   vkCmdBindIndexBuffer(commandBuffers[currentImage], info.indexBuffer, 0, VK_INDEX_TYPE_UINT16);

   vkCmdDrawIndexed(commandBuffers[currentImage], static_cast<uint32_t>(info.numIndices), 1, 0, 0, 0);
   */
}

VkShaderModule GraphicsPipeline_Vulkan::createShaderModule(const vector<char>& code) {
   VkShaderModuleCreateInfo createInfo = {};
   createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
   createInfo.codeSize = code.size();
   createInfo.pCode = reinterpret_cast<const uint32_t*>(code.data());

   VkShaderModule shaderModule;
   if (vkCreateShaderModule(this->device, &createInfo, nullptr, &shaderModule) != VK_SUCCESS) {
      throw runtime_error("failed to create shader module!");
   }

   return shaderModule;
}

vector<char> GraphicsPipeline_Vulkan::readFile(const string& filename) {
   ifstream file(filename, ios::ate | ios::binary);

   if (!file.is_open()) {
      throw runtime_error("failed to open file!");
   }

   size_t fileSize = (size_t)file.tellg();
   vector<char> buffer(fileSize);

   file.seekg(0);
   file.read(buffer.data(), fileSize);

   file.close();

   return buffer;
}

void GraphicsPipeline_Vulkan::cleanup() {
   vkDestroyPipeline(device, pipeline, nullptr);
   vkDestroyDescriptorPool(device, descriptorPool, nullptr);
   vkDestroyPipelineLayout(device, pipelineLayout, nullptr);
}

void GraphicsPipeline_Vulkan::cleanupBuffers() {
   vkDestroyDescriptorSetLayout(device, descriptorSetLayout, nullptr);

   vkDestroyBuffer(device, vertexBuffer, nullptr);
   vkFreeMemory(device, vertexBufferMemory, nullptr);
   vkDestroyBuffer(device, indexBuffer, nullptr);
   vkFreeMemory(device, indexBufferMemory, nullptr);
}