/*
 * Copyright 2022 Frederico de Oliveira Linhares
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "graphics_pipeline_2d.hpp"

#include <array>

#include "../core.hpp"
#include "sprite.hpp"

namespace
{

void
load_projection_uniform_buffer(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  try
  {
    self->ub_projection.reserve(cg_core.vk_swapchain->images_count);
    for(auto i{0}; i < cg_core.vk_swapchain->images_count; i++)
      self->ub_projection.emplace_back(
	cg_core.vk_device_with_swapchain, sizeof(VK::UBOProjection));
  }
  catch(const std::exception& e)
  {
    throw CommandError{e.what()};
  }
}

void
unload_projection_uniform_buffer(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  self->ub_projection.clear();
}

void
load_descriptor_pool(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  VkDescriptorPoolSize descriptor_pool_size{};
  descriptor_pool_size.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  descriptor_pool_size.descriptorCount = self->ub_projection.size();

  VkDescriptorPoolCreateInfo pool_info{};
  pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  pool_info.pNext = nullptr;
  pool_info.flags = 0;
  pool_info.maxSets = self->ub_projection.size();
  pool_info.poolSizeCount = 1;
  pool_info.pPoolSizes = &descriptor_pool_size;

  if(vkCreateDescriptorPool(
       cg_core.vk_device_with_swapchain->device, &pool_info, nullptr,
       &self->descriptor_pool) != VK_SUCCESS)
    throw CommandError{"Failed to create a Vulkan descriptor pool."};
}

void
unload_descriptor_pool(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  vkDestroyDescriptorPool(
    cg_core.vk_device_with_swapchain->device, self->descriptor_pool,
    nullptr);
}

void
load_projection_descriptor_sets(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  std::vector<VkDescriptorSetLayout> layouts(
    self->ub_projection.size(),
    cg_core.vk_graphics_pipeline_2d_layout->descriptor_set_projection);

  VkDescriptorSetAllocateInfo alloc_info{};
  alloc_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
  alloc_info.descriptorPool = self->descriptor_pool;
  alloc_info.descriptorSetCount = layouts.size();
  alloc_info.pSetLayouts = layouts.data();

  self->projection_descriptor_sets.resize(layouts.size());
  if(vkAllocateDescriptorSets(
       cg_core.vk_device_with_swapchain->device, &alloc_info,
       self->projection_descriptor_sets.data()) != VK_SUCCESS)
    throw CommandError{"Failed to create Vulkan projection descriptor sets."};
}

void
load_resources_to_descriptor_sets(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  for(auto i{0}; i < self->ub_projection.size(); i++)
  {
    VkDescriptorBufferInfo projection_info{};
    projection_info.buffer = self->ub_projection[i].buffer;
    projection_info.offset = 0;
    projection_info.range = sizeof(VK::UBOProjection);

    std::array<VkWriteDescriptorSet, 1> write_descriptors{};
    write_descriptors[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptors[0].dstSet = self->projection_descriptor_sets[i];
    write_descriptors[0].dstBinding = 0;
    write_descriptors[0].dstArrayElement = 0;
    write_descriptors[0].descriptorCount = 1;
    write_descriptors[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    write_descriptors[0].pBufferInfo = &projection_info;
    write_descriptors[0].pImageInfo = nullptr;
    write_descriptors[0].pTexelBufferView = nullptr;

    vkUpdateDescriptorSets(
      cg_core.vk_device_with_swapchain->device, write_descriptors.size(),
      write_descriptors.data(), 0, nullptr);

    VK::UBOProjection ubo_projection;
    ubo_projection.proj = glm::ortho(
      0.0f, static_cast<float>(cg_core.game_width),
      0.0f, static_cast<float>(cg_core.game_height),
      0.0f, 100.0f);
    self->ub_projection[i].copy_data(&ubo_projection);
  }
}

void
load_framebuffer(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  self->swapchain_framebuffers.resize(cg_core.vk_swapchain->images_count);
  for (auto i{0}; i < cg_core.vk_swapchain->images_count; i++)
  {
    std::array<VkImageView, 1> attachments = {
      cg_core.vk_swapchain->image_views[i]
    };

    VkFramebufferCreateInfo framebuffer_info{};
    framebuffer_info.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
    framebuffer_info.renderPass =
      cg_core.vk_graphics_pipeline_2d_layout->render_pass;
    framebuffer_info.attachmentCount = attachments.size();
    framebuffer_info.pAttachments = attachments.data();
    framebuffer_info.width = cg_core.screen_width;
    framebuffer_info.height = cg_core.screen_height;
    framebuffer_info.layers = 1;

    if(vkCreateFramebuffer(
	 cg_core.vk_device_with_swapchain->device, &framebuffer_info, nullptr,
	 &self->swapchain_framebuffers[i])
       != VK_SUCCESS)
      throw CommandError{"Failed to create Vulkan Framebuffer."};
  }
}

void
unload_framebuffer(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  for(auto framebuffer: self->swapchain_framebuffers)
    vkDestroyFramebuffer(
      cg_core.vk_device_with_swapchain->device, framebuffer, nullptr);
}

void
load_pipeline(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  VkPipelineShaderStageCreateInfo vert_shader_stage_info{};
  vert_shader_stage_info.sType =
      VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  vert_shader_stage_info.pNext = nullptr;
  vert_shader_stage_info.flags = 0;
  vert_shader_stage_info.stage = VK_SHADER_STAGE_VERTEX_BIT;
  vert_shader_stage_info.module =
    cg_core.vk_device_with_swapchain->vert2d_shader_module;
  vert_shader_stage_info.pName = "main";
  vert_shader_stage_info.pSpecializationInfo = nullptr;

  VkPipelineShaderStageCreateInfo frag_shader_stage_info{};
  frag_shader_stage_info.sType =
      VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  frag_shader_stage_info.pNext = nullptr;
  frag_shader_stage_info.flags = 0;
  frag_shader_stage_info.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
  frag_shader_stage_info.module =
    cg_core.vk_device_with_swapchain->frag2d_shader_module;
  frag_shader_stage_info.pName = "main";
  frag_shader_stage_info.pSpecializationInfo = nullptr;

  VkPipelineShaderStageCreateInfo shader_stages[] = {
    vert_shader_stage_info,
    frag_shader_stage_info
  };

  VkVertexInputBindingDescription vertex_input_binding{};
  vertex_input_binding.binding = 0;
  vertex_input_binding.stride = sizeof(glm::vec2);
  vertex_input_binding.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;

  std::array<VkVertexInputAttributeDescription, 1> vertex_attribute{};
  // Texture coordinate.
  vertex_attribute[0].location = 0;
  vertex_attribute[0].binding = 0;
  vertex_attribute[0].format = VK_FORMAT_R32G32_SFLOAT;
  vertex_attribute[0].offset = 0;

  VkPipelineVertexInputStateCreateInfo vertex_input_info = {};
  vertex_input_info.sType =
      VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
  vertex_input_info.pNext = nullptr;
  vertex_input_info.flags = 0;
  vertex_input_info.vertexBindingDescriptionCount = 1;
  vertex_input_info.pVertexBindingDescriptions = &vertex_input_binding;
  vertex_input_info.vertexAttributeDescriptionCount =
      static_cast<uint32_t>(vertex_attribute.size());
  vertex_input_info.pVertexAttributeDescriptions = vertex_attribute.data();

  VkPipelineInputAssemblyStateCreateInfo input_assembly = {};
  input_assembly.sType =
      VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
  input_assembly.pNext = nullptr;
  input_assembly.flags = 0;
  input_assembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_STRIP;
  input_assembly.primitiveRestartEnable = VK_FALSE;

  VkViewport viewport = {};
  viewport.x = 0.0f;
  viewport.y = 0.0f;
  viewport.width = static_cast<float>(cg_core.screen_width);
  viewport.height = static_cast<float>(cg_core.screen_height);
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;

  VkRect2D scissor = {};
  scissor.offset = {0, 0};
  scissor.extent = {cg_core.screen_width, cg_core.screen_height};

  VkPipelineViewportStateCreateInfo viewport_state = {};
  viewport_state.sType =
      VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
  viewport_state.pNext = nullptr;
  viewport_state.flags = 0;
  viewport_state.viewportCount = 1;
  viewport_state.pViewports = &viewport;
  viewport_state.scissorCount = 1;
  viewport_state.pScissors = &scissor;

  VkPipelineRasterizationStateCreateInfo rasterizer = {};
  rasterizer.sType =
      VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
  rasterizer.pNext = nullptr;
  rasterizer.flags = 0;
  rasterizer.depthClampEnable = VK_FALSE;
  rasterizer.rasterizerDiscardEnable = VK_FALSE;
  rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
  rasterizer.cullMode = VK_CULL_MODE_NONE;
  rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
  rasterizer.depthBiasEnable = VK_FALSE;
  rasterizer.depthBiasConstantFactor = 0.0f;
  rasterizer.depthBiasClamp = 0.0f;
  rasterizer.depthBiasSlopeFactor = 0.0f;
  rasterizer.lineWidth = 1.0f;

  VkPipelineMultisampleStateCreateInfo multisampling = {};
  multisampling.sType =
      VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
  multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
  multisampling.sampleShadingEnable = VK_FALSE;
  multisampling.minSampleShading = 1.0f;
  multisampling.pSampleMask = nullptr;
  multisampling.alphaToCoverageEnable = VK_FALSE;
  multisampling.alphaToOneEnable = VK_FALSE;

  VkPipelineColorBlendAttachmentState color_blend_attachment = {};
  color_blend_attachment.blendEnable = VK_FALSE;
  color_blend_attachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
  color_blend_attachment.dstColorBlendFactor = VK_BLEND_FACTOR_ZERO;
  color_blend_attachment.colorBlendOp = VK_BLEND_OP_ADD;
  color_blend_attachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
  color_blend_attachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
  color_blend_attachment.alphaBlendOp = VK_BLEND_OP_ADD;
  color_blend_attachment.colorWriteMask =
      VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT |
      VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;

  VkPipelineColorBlendStateCreateInfo color_blending = {};
  color_blending.sType =
      VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
  color_blending.pNext = nullptr;
  color_blending.flags = 0;
  color_blending.logicOpEnable = VK_FALSE;
  color_blending.logicOp = VK_LOGIC_OP_COPY;
  color_blending.attachmentCount = 1;
  color_blending.pAttachments = &color_blend_attachment;
  color_blending.blendConstants[0] = 0.0f;
  color_blending.blendConstants[1] = 0.0f;
  color_blending.blendConstants[2] = 0.0f;
  color_blending.blendConstants[3] = 0.0f;

  VkDynamicState dynamic_states[] = {
    VK_DYNAMIC_STATE_VIEWPORT,
    VK_DYNAMIC_STATE_LINE_WIDTH
  };

  VkPipelineDynamicStateCreateInfo dynamic_state_info = {};
  dynamic_state_info.sType =
      VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
  dynamic_state_info.dynamicStateCount = 2;
  dynamic_state_info.pDynamicStates = dynamic_states;

  VkGraphicsPipelineCreateInfo pipeline_info{};
  pipeline_info.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
  pipeline_info.pNext = nullptr;
  pipeline_info.flags = 0;
  pipeline_info.stageCount = 2;
  pipeline_info.pStages = shader_stages;
  pipeline_info.pVertexInputState = &vertex_input_info;
  pipeline_info.pInputAssemblyState = &input_assembly;
  pipeline_info.pTessellationState = nullptr;
  pipeline_info.pViewportState = &viewport_state;
  pipeline_info.pRasterizationState = &rasterizer;
  pipeline_info.pMultisampleState = &multisampling;
  pipeline_info.pDepthStencilState = nullptr;
  pipeline_info.pColorBlendState = &color_blending;
  pipeline_info.pDynamicState = &dynamic_state_info;
  pipeline_info.layout = cg_core.vk_graphics_pipeline_2d_layout->pipeline;
  pipeline_info.renderPass =
    cg_core.vk_graphics_pipeline_2d_layout->render_pass;
  pipeline_info.subpass = 0;
  pipeline_info.basePipelineHandle = VK_NULL_HANDLE;
  pipeline_info.basePipelineIndex = -1;

  if(vkCreateGraphicsPipelines(
       cg_core.vk_device_with_swapchain->device, VK_NULL_HANDLE, 1,
       &pipeline_info, nullptr, &self->graphic_pipeline)
     != VK_SUCCESS)
    throw CommandError{"Failed to create graphics pipeline."};
}

void
unload_pipeline(void *obj)
{
  auto self = static_cast<VK::GraphicsPipeline2D*>(obj);

  vkDestroyPipeline(
    cg_core.vk_device_with_swapchain->device, self->graphic_pipeline, nullptr);
}

const CommandChain loader{
  {&load_projection_uniform_buffer, &unload_projection_uniform_buffer},
  {&load_descriptor_pool, &unload_descriptor_pool},
  // By destroying the pool the sets are also destroyed.
  {&load_projection_descriptor_sets, nullptr},
  {&load_resources_to_descriptor_sets, nullptr},
  {&load_framebuffer, &unload_framebuffer},
  {&load_pipeline, &unload_pipeline}
};

}

namespace VK
{

GraphicsPipeline2D::GraphicsPipeline2D():
  sprites_to_draw{cg_core.vk_swapchain->images_count}
{
  loader.execute(this);
}

GraphicsPipeline2D::~GraphicsPipeline2D()
{
  loader.revert(this);
}

void
GraphicsPipeline2D::draw(
  const VkCommandBuffer draw_command_buffer, const size_t current_frame,
  const size_t next_frame, const uint32_t image_index)
{
  // Load command
  {
    VkRenderPassBeginInfo render_pass_begin{};
    render_pass_begin.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    render_pass_begin.pNext = nullptr;
    render_pass_begin.renderPass =
      cg_core.vk_graphics_pipeline_2d_layout->render_pass;
    render_pass_begin.framebuffer = this->swapchain_framebuffers[image_index];
    render_pass_begin.renderArea.offset = {0, 0};
    render_pass_begin.renderArea.extent = {
      cg_core.screen_width, cg_core.screen_height};
    render_pass_begin.clearValueCount = 0;
    render_pass_begin.pClearValues = nullptr;

    vkCmdBeginRenderPass(
      draw_command_buffer, &render_pass_begin, VK_SUBPASS_CONTENTS_INLINE);

    VkViewport vk_viewport{};
    vk_viewport.width = static_cast<float>(cg_core.screen_width);
    vk_viewport.height = static_cast<float>(cg_core.screen_height);
    vk_viewport.minDepth = 0.0f;
    vk_viewport.maxDepth = 1.0f;
    vkCmdSetViewport(draw_command_buffer, 0, 1, &vk_viewport);

    VkRect2D vk_scissor{};
    vk_scissor.extent.width = cg_core.screen_width;
    vk_scissor.extent.height = cg_core.screen_height;
    vk_scissor.offset.x = 0;
    vk_scissor.offset.y = 0;
    vkCmdSetScissor(draw_command_buffer, 0, 1, &vk_scissor);

    // Draw sprites
    for(auto& [sprite, positions]: this->sprites_to_draw[current_frame])
    {
      // Commands
      {
	std::array<VkDescriptorSet, 2> vk_descriptor_sets{
	  this->projection_descriptor_sets[image_index],
	  sprite->descriptor_sets[image_index]};
	VkDeviceSize offsets[]{0};

	vkCmdBindDescriptorSets(
	  draw_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
	  cg_core.vk_graphics_pipeline_2d_layout->pipeline, 0,
	  vk_descriptor_sets.size(), vk_descriptor_sets.data(), 0, nullptr);
	vkCmdBindPipeline(
	  draw_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
	  this->graphic_pipeline);
	vkCmdBindVertexBuffers(
	  draw_command_buffer, 0, 1, &sprite->vertex_buffer->buffer, offsets);
	vkCmdDraw(
	  draw_command_buffer, Sprite::vertex_count, positions.size(), 0, 0);
      }

      VK::UBOSpritePositions ubo_sprite_positions;
      for(auto i{0}; i < positions.size(); i++)
	ubo_sprite_positions.positions[i] = positions[i];
      sprite->ub_sprite_positions[image_index].copy_data(
	&ubo_sprite_positions);
    }
  }
  vkCmdEndRenderPass(draw_command_buffer);

  // Prepare for the next frame.
  this->sprites_to_draw[next_frame].clear();
}

}