#include "VulkanRenderer.h"
#include "VertexUniformBufferObject.h"
#include <complex>

VulkanRenderer::VulkanRenderer(QVulkanWindow *window, WorldView *worldView)
    : m_window(window), m_worldView(worldView) {}

void VulkanRenderer::createPipeline() {
  VkResult result;

  // Note: switched to push constants instead
  //  VkDescriptorSetLayoutBinding uboLayoutBinding{};
  //  uboLayoutBinding.binding = 0;
  //  uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  //  uboLayoutBinding.descriptorCount = 1;
  //  uboLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;
  //
  //  VkDescriptorSetLayoutCreateInfo layoutInfo{};
  //  layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  //  layoutInfo.bindingCount = 1;
  //  layoutInfo.pBindings = &uboLayoutBinding;
  //
  //  result = m_devFuncs->vkCreateDescriptorSetLayout(
  //      m_window->device(), &layoutInfo, VK_NULL_HANDLE,
  //      &m_descriptorSetLayout);
  //  if (result != VK_SUCCESS) {
  //    qFatal("Failed to create descriptor set layout, code %d", result);
  //  }

  VkPushConstantRange pushConstant{};
  pushConstant.offset = 0;
  // might be too big for some graphics cards, spec requires min 128 bytes (the
  // validation layer will tell us)
  pushConstant.size = sizeof(VertexUniformBufferObject);
  pushConstant.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;

  auto shaders =
      std::array<std::string, 2>{"shaders/vert.spv", "shaders/frag.spv"};
  auto loadedShaders =
      m_shaderLoader.loadShaders(m_window->device(), m_devFuncs, shaders);
  VkPipelineShaderStageCreateInfo vertShaderStageInfo{};
  vertShaderStageInfo.sType =
      VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
  vertShaderStageInfo.module = loadedShaders[0];
  vertShaderStageInfo.pName = "main";
  VkPipelineShaderStageCreateInfo fragShaderStageInfo{};
  fragShaderStageInfo.sType =
      VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
  fragShaderStageInfo.module = loadedShaders[1];
  fragShaderStageInfo.pName = "main";
  std::array<VkPipelineShaderStageCreateInfo, 2> shaderStages{
      vertShaderStageInfo, fragShaderStageInfo};

  VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
  pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
  // Note: switched to push constants instead
  //  pipelineLayoutInfo.setLayoutCount = 1;
  //  pipelineLayoutInfo.pSetLayouts = &m_descriptorSetLayout;
  pipelineLayoutInfo.pushConstantRangeCount = 1;
  pipelineLayoutInfo.pPushConstantRanges = &pushConstant;

  result = m_devFuncs->vkCreatePipelineLayout(
      m_window->device(), &pipelineLayoutInfo, VK_NULL_HANDLE,
      &m_pipelineLayout);
  if (result != VK_SUCCESS) {
    qFatal("Failed to create pipeline layout, code: %d", result);
  }

  VkPipelineVertexInputStateCreateInfo vertexInputInfo{};
  vertexInputInfo.sType =
      VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
  vertexInputInfo.vertexBindingDescriptionCount = 0;
  vertexInputInfo.pVertexBindingDescriptions = VK_NULL_HANDLE; // Optional
  vertexInputInfo.vertexAttributeDescriptionCount = 0;
  vertexInputInfo.pVertexAttributeDescriptions = VK_NULL_HANDLE; // Optional

  VkGraphicsPipelineCreateInfo pipelineInfo{};
  pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
  pipelineInfo.stageCount = shaderStages.size();
  pipelineInfo.pStages = shaderStages.data();
  pipelineInfo.pVertexInputState = &vertexInputInfo;

  VkPipelineInputAssemblyStateCreateInfo ia;
  memset(&ia, 0, sizeof(ia));
  ia.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
  ia.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
  pipelineInfo.pInputAssemblyState = &ia;

  VkPipelineViewportStateCreateInfo vp;
  memset(&vp, 0, sizeof(vp));
  vp.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
  vp.viewportCount = 1;
  vp.scissorCount = 1;
  pipelineInfo.pViewportState = &vp;

  VkPipelineRasterizationStateCreateInfo rs;
  memset(&rs, 0, sizeof(rs));
  rs.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
  rs.polygonMode = VK_POLYGON_MODE_FILL;
  //  rs.cullMode = VK_CULL_MODE_BACK_BIT;
  rs.cullMode = VK_CULL_MODE_NONE;
  rs.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
  rs.lineWidth = 1.0f;
  pipelineInfo.pRasterizationState = &rs;

  VkPipelineMultisampleStateCreateInfo ms;
  memset(&ms, 0, sizeof(ms));
  ms.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
  ms.rasterizationSamples = m_window->sampleCountFlagBits();
  pipelineInfo.pMultisampleState = &ms;

  VkPipelineDepthStencilStateCreateInfo ds;
  memset(&ds, 0, sizeof(ds));
  ds.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
  ds.depthTestEnable = VK_FALSE;
  ds.depthWriteEnable = VK_TRUE;
  ds.depthCompareOp = VK_COMPARE_OP_LESS_OR_EQUAL;
  pipelineInfo.pDepthStencilState = &ds;

  VkPipelineColorBlendStateCreateInfo cb;
  memset(&cb, 0, sizeof(cb));
  cb.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
  VkPipelineColorBlendAttachmentState att;
  memset(&att, 0, sizeof(att));
  att.colorWriteMask = 0xF;
  cb.attachmentCount = 1;
  cb.pAttachments = &att;
  pipelineInfo.pColorBlendState = &cb;

  VkDynamicState dynamicStates[] = {VK_DYNAMIC_STATE_VIEWPORT,
                                    VK_DYNAMIC_STATE_SCISSOR};
  VkPipelineDynamicStateCreateInfo dyn;
  memset(&dyn, 0, sizeof(dyn));
  dyn.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
  dyn.dynamicStateCount = sizeof(dynamicStates) / sizeof(VkDynamicState);
  dyn.pDynamicStates = dynamicStates;
  pipelineInfo.pDynamicState = &dyn;

  pipelineInfo.layout = m_pipelineLayout;
  pipelineInfo.renderPass = m_window->defaultRenderPass();
  pipelineInfo.basePipelineHandle = VK_NULL_HANDLE; // Optional
  pipelineInfo.basePipelineIndex = -1;              // Optional

  result = m_devFuncs->vkCreateGraphicsPipelines(
      m_window->device(), VK_NULL_HANDLE, 1, &pipelineInfo, VK_NULL_HANDLE,
      &m_graphicsPipeline);
  if (result != VK_SUCCESS) {
    qFatal("Failed to create graphics pipeline: code %d", result);
  }
}

void VulkanRenderer::initResources() {
  QVulkanWindowRenderer::initResources();
  m_devFuncs = m_window->vulkanInstance()->deviceFunctions(m_window->device());
  createPipeline();
}

void VulkanRenderer::releaseResources() {
  QVulkanWindowRenderer::releaseResources();
  m_devFuncs->vkDestroyPipeline(m_window->device(), m_graphicsPipeline,
                                VK_NULL_HANDLE);
  m_devFuncs->vkDestroyPipelineLayout(m_window->device(), m_pipelineLayout,
                                      VK_NULL_HANDLE);
  //  m_devFuncs->vkDestroyDescriptorSetLayout(
  //      m_window->device(), m_descriptorSetLayout, VK_NULL_HANDLE);
  m_shaderLoader.destroyShaders(m_window->device(), m_devFuncs);
}

void VulkanRenderer::initSwapChainResources() {
  QVulkanWindowRenderer::initSwapChainResources();
  //  m_uniformBuffers.createBuffers(m_window, m_devFuncs);
  m_worldView->initializeProjectionMatrix(m_window);
}

void VulkanRenderer::releaseSwapChainResources() {
  QVulkanWindowRenderer::releaseSwapChainResources();
  //  m_uniformBuffers.destroyBuffers(m_window, m_devFuncs);
}

void VulkanRenderer::startNextFrame() {
  VkCommandBuffer cmdBuf = m_window->currentCommandBuffer();

  VkClearDepthStencilValue clearDS = {1, 0};
  std::array<VkClearValue, 3> clearValue{};
  float flash =
      std::abs(std::sin(static_cast<float>(m_frameNumber) / 80.f)) / 2.f;
  clearValue[0].color = {{0.0f, 0.0f, flash, 1.0f}};
  clearValue[1].depthStencil = clearDS;
  clearValue[2].color = clearValue[0].color;

  VkRenderPassBeginInfo rpInfo{};
  rpInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
  rpInfo.pNext = VK_NULL_HANDLE;

  rpInfo.renderPass = m_window->defaultRenderPass();
  rpInfo.renderArea.offset.x = 0;
  rpInfo.renderArea.offset.y = 0;
  rpInfo.renderArea.extent.width = m_window->width();
  rpInfo.renderArea.extent.height = m_window->height();
  rpInfo.clearValueCount =
      m_window->sampleCountFlagBits() > VK_SAMPLE_COUNT_1_BIT ? 3 : 2;
  rpInfo.pClearValues = clearValue.data();
  rpInfo.framebuffer = m_window->currentFramebuffer();

  m_devFuncs->vkCmdBeginRenderPass(cmdBuf, &rpInfo, VK_SUBPASS_CONTENTS_INLINE);

  const QSize renderTargetSize = m_window->swapChainImageSize();
  VkViewport viewport = {
      0, 0, float(renderTargetSize.width()), float(renderTargetSize.height()),
      0, 1};
  m_devFuncs->vkCmdSetViewport(cmdBuf, 0, 1, &viewport);

  VkRect2D scissor = {{0, 0},
                      {uint32_t(renderTargetSize.width()),
                       uint32_t(renderTargetSize.height())}};
  m_devFuncs->vkCmdSetScissor(cmdBuf, 0, 1, &scissor);

  m_devFuncs->vkCmdBindPipeline(cmdBuf, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                m_graphicsPipeline);

  const VertexUniformBufferObject &mvp = m_worldView->getBufferObject();
  m_devFuncs->vkCmdPushConstants(cmdBuf, m_pipelineLayout,
                                 VK_SHADER_STAGE_VERTEX_BIT, 0,
                                 sizeof(VertexUniformBufferObject), &mvp);

  m_devFuncs->vkCmdDraw(cmdBuf, 12, 1, 0, 0);

  m_devFuncs->vkCmdEndRenderPass(cmdBuf);
  m_window->frameReady();
  ++m_frameNumber;
  m_window->requestUpdate();
}