/*
 * Copyright 2022-2024 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 "renderer.hpp"

#include <array>

#include "core.hpp"
#include "uniform_data_object.hpp"

namespace
{

void
load_descriptor_pool(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  uint32_t uniform_buffer_count = 0;
  for(auto &view : self->views_3d)
    uniform_buffer_count += (view->ub_3d.size() + view->ub_2d.size());
  for(auto &view : self->views_2d)
    uniform_buffer_count += (view->ub_2d.size());

  VkDescriptorPoolSize descriptor_pool_size{};
  descriptor_pool_size.type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  descriptor_pool_size.descriptorCount = uniform_buffer_count;

  VkDescriptorPoolCreateInfo pool_info{};
  pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  pool_info.pNext = nullptr;
  pool_info.flags = 0;
  pool_info.maxSets = uniform_buffer_count;
  pool_info.poolSizeCount = 1;
  pool_info.pPoolSizes = &descriptor_pool_size;

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

  for(auto &view : self->views_3d)
    view->load_descriptor_sets(self->descriptor_pool);
  for(auto &view : self->views_2d)
    view->load_descriptor_sets(self->descriptor_pool);
}

void
unload_descriptor_pool(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  for(auto &view : self->views_3d) view->unload_descriptor_sets();
  for(auto &view : self->views_2d) view->unload_descriptor_sets();

  vkDestroyDescriptorPool(
    BluCat::core.vk_device_with_swapchain->device, self->descriptor_pool,
    nullptr);
}

void
load_queue_family(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  self->queue_family =
    BluCat::core.vk_device_with_swapchain->get_queue_family_with_presentation();
}

void
load_command_pool(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  VkCommandPoolCreateInfo create_info{};
  create_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
  create_info.pNext = nullptr;
  create_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
  create_info.queueFamilyIndex = self->queue_family->family_index;

  vkCreateCommandPool(
    self->queue_family->device->device, &create_info, nullptr,
    &self->command_pool);
}

void
unload_command_pool(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  vkWaitForFences(BluCat::core.vk_device_with_swapchain->device,
		  BluCat::Swapchain::max_frames_in_flight,
		  BluCat::core.vk_swapchain->in_flight_fences.data(), VK_TRUE,
		  std::numeric_limits<uint64_t>::max());
  vkDestroyCommandPool(
    self->queue_family->device->device, self->command_pool, nullptr);
}

void
load_draw_command_buffer(void *obj)
{
  auto self = static_cast<BluCat::Renderer*>(obj);

  // FIXME: 3 is a magical number, triple buffering.
  self->draw_command_buffers.resize(3);

  VkCommandBufferAllocateInfo allocate_info{};
  allocate_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocate_info.pNext = nullptr;
  allocate_info.commandPool = self->command_pool;
  allocate_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
  allocate_info.commandBufferCount = self->draw_command_buffers.size();

  if(vkAllocateCommandBuffers(
       self->queue_family->device->device, &allocate_info,
       self->draw_command_buffers.data()) != VK_SUCCESS)
    throw CommandError{"Vulkan draw command buffers could not be allocated."};
}

const CommandChain loader{
  {&load_descriptor_pool, &unload_descriptor_pool},
  {&load_queue_family, nullptr},
  {&load_command_pool, &unload_command_pool},
  {&load_draw_command_buffer, nullptr}
};

}

namespace BluCat
{

Renderer::Renderer(std::vector<std::shared_ptr<View2D>> views_2d,
		   std::vector<std::shared_ptr<View3D>> views_3d):
  skeletal_models_to_draw{BluCat::core.vk_swapchain->images_count},
  static_models_to_draw{BluCat::core.vk_swapchain->images_count},
  sprites_3d_to_draw{BluCat::core.vk_swapchain->images_count},
  views_2d{views_2d},
  views_3d{views_3d}
{
  loader.execute(this);
}

Renderer::Renderer(std::initializer_list<std::shared_ptr<View2D>> views_2d,
		   std::initializer_list<std::shared_ptr<View3D>> views_3d):
  Renderer(std::vector(views_2d), std::vector(views_3d))
{
}

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

void
Renderer::draw()
{
  auto fence_status = vkGetFenceStatus(
    BluCat::core.vk_device_with_swapchain->device,
    BluCat::core.vk_swapchain->in_flight_fences[
      BluCat::core.vk_swapchain->current_frame]);

  if(fence_status == VK_SUCCESS)
  {
    auto next_frame = BluCat::core.vk_swapchain->current_frame + 1;
    if(next_frame == Swapchain::max_frames_in_flight) next_frame = 0;

    vkResetFences(core.vk_device_with_swapchain->device, 1,
		  &BluCat::core.vk_swapchain->in_flight_fences[
		    BluCat::core.vk_swapchain->current_frame]);

    uint32_t image_index;
    vkAcquireNextImageKHR(
      BluCat::core.vk_device_with_swapchain->device,
      BluCat::core.vk_swapchain->swapchain,
			std::numeric_limits<uint64_t>::max(),
      BluCat::core.vk_swapchain->image_available_semaphores[
	BluCat::core.vk_swapchain->current_frame], VK_NULL_HANDLE, &image_index);

    VkCommandBuffer draw_command_buffer =
      this->draw_command_buffers[BluCat::core.vk_swapchain->current_frame];
    vkResetCommandBuffer(draw_command_buffer, 0);

    // Begin command buffer.
    {
      VkCommandBufferBeginInfo begin_info{};
      begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
      begin_info.flags = 0;
      begin_info.pInheritanceInfo = nullptr;
      if (vkBeginCommandBuffer(draw_command_buffer, &begin_info) != VK_SUCCESS)
	throw std::runtime_error{"Failed to beggin draw command buffer."};
    }

    // 3D drawing.
    {
      // Dark gray blue.
      std::array<VkClearValue, 2> clear_values{};
      clear_values[0].color = {0.12f, 0.12f, 0.18f, 1.0f};
      clear_values[1].depthStencil = {1.0f, 0};

      { // Update world uniform buffer
	UDOWorld3D_Vert ubo_world_3d_vert{};
	ubo_world_3d_vert.ambient_light_color =
	  glm::vec4{0.25, 0.25, 0.25, 1.0};
	BluCat::core.vk_light->ub_world_vert[image_index].copy_data(
	  &ubo_world_3d_vert);

	UDOWorld3D_Frag ubo_world_3d_frag{};
	ubo_world_3d_frag.directional_light_direction =
	  glm::vec3{-0.57735, 0.57735, -0.57735};
	ubo_world_3d_frag.directional_light_color =
	  glm::vec4{0.8, 0.8, 0.8, 1.0};
	BluCat::core.vk_light->ub_world_frag[image_index].copy_data(
	  &ubo_world_3d_frag);
      }

      VkRenderPassBeginInfo render_pass_begin{};
      render_pass_begin.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
      render_pass_begin.pNext = nullptr;
      render_pass_begin.renderPass = BluCat::core.vk_render_pass->pipeline_3d;
      render_pass_begin.framebuffer =
	BluCat::core.vk_framebuffer->pipeline_3d[image_index];
      render_pass_begin.renderArea.offset = {0, 0};
      render_pass_begin.renderArea.extent = {
	static_cast<uint32_t>(BluCat::core.display_width),
	static_cast<uint32_t>(BluCat::core.display_height)};
      render_pass_begin.clearValueCount = clear_values.size();
      render_pass_begin.pClearValues = clear_values.data();

      vkCmdBeginRenderPass(
	draw_command_buffer, &render_pass_begin, VK_SUBPASS_CONTENTS_INLINE);
    }

		for(auto &view: this->views_3d)
		{
			{ // Set viewport
				VkViewport vk_viewport{};
				vk_viewport.x = view->region.x;
				vk_viewport.y = view->region.y;
				vk_viewport.width = view->region.z;
				vk_viewport.height = view->region.w;
				vk_viewport.minDepth = 0.0f;
				vk_viewport.maxDepth = 1.0f;
				vkCmdSetViewport(draw_command_buffer, 0, 1, &vk_viewport);

				VkRect2D vk_scissor{};
				vk_scissor.offset.x = static_cast<int32_t>(view->region.x);
				vk_scissor.offset.y = static_cast<int32_t>(view->region.y);
				vk_scissor.extent.width = static_cast<uint32_t>(view->region.z);
				vk_scissor.extent.height = static_cast<uint32_t>(view->region.w);
				vkCmdSetScissor(draw_command_buffer, 0, 1, &vk_scissor);
			}

			BluCat::core.vk_graphics_pipeline_3d->draw(
				view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
				image_index);

			BluCat::core.vk_graphics_pipeline_sprite_3d->draw(
				view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
				image_index);

			BluCat::core.vk_graphics_pipeline_3d_skeletal->draw(
				view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
				image_index);

			{ // Update view uniform buffers
				BluCat::UDOView3D ubo_view_3d{};

				// View matrix.
				glm::mat4 translation_matrix{1.0f};
				translation_matrix = glm::translate(
					translation_matrix, *view->camera_position);
				glm::mat4 rotation_matrix{glm::toMat4(*view->camera_orientation)};
				ubo_view_3d.view = glm::inverse(translation_matrix * rotation_matrix);

				// Projection matrix.
				ubo_view_3d.proj = glm::perspective(
					glm::radians(view->field_of_view),
					view->region.z / view->region.w,
					0.1f, 100.0f);

				view->ub_3d[image_index].copy_data(&ubo_view_3d);
			}
		}

    vkCmdEndRenderPass(draw_command_buffer);

    { // 2D render pass
      VkRenderPassBeginInfo render_pass_begin{};
      render_pass_begin.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
      render_pass_begin.pNext = nullptr;
      render_pass_begin.renderPass = BluCat::core.vk_render_pass->pipeline_2d;
      render_pass_begin.framebuffer =
	BluCat::core.vk_framebuffer->pipeline_2d[image_index];
      render_pass_begin.renderArea.offset = {0, 0};
      render_pass_begin.renderArea.extent = {
	static_cast<uint32_t>(BluCat::core.display_width),
	static_cast<uint32_t>(BluCat::core.display_height)};
      render_pass_begin.clearValueCount = 0;
      render_pass_begin.pClearValues = nullptr;

      vkCmdBeginRenderPass(
	draw_command_buffer, &render_pass_begin, VK_SUBPASS_CONTENTS_INLINE);

    }

    { // 2D solid drawing
      for(auto &view: this->views_2d)
	BluCat::core.vk_graphics_pipeline_2d_solid->draw(
	  view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
	  next_frame, image_index);

      for(auto &view: this->views_3d)
	BluCat::core.vk_graphics_pipeline_2d_solid->draw(
	  view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
	  next_frame, image_index);
    }

    { // 2D wired drawing
      for(auto &view: this->views_2d)
	BluCat::core.vk_graphics_pipeline_2d_wired->draw(
	  view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
	  next_frame, image_index);

      for(auto &view: this->views_3d)
	BluCat::core.vk_graphics_pipeline_2d_wired->draw(
	  view, draw_command_buffer, BluCat::core.vk_swapchain->current_frame,
	  next_frame, image_index);
    }

    vkCmdEndRenderPass(draw_command_buffer);

    // End command buffer.
    if(vkEndCommandBuffer(draw_command_buffer) != VK_SUCCESS)
      throw std::runtime_error{"Failed to end draw command buffer."};

    // Submit drawing command.
    {
      auto queue{this->queue_family->get_queue()};

      VkSemaphore wait_semaphores[]{
	BluCat::core.vk_swapchain->image_available_semaphores[
	  BluCat::core.vk_swapchain->current_frame]};
      VkPipelineStageFlags wait_stages[] =
	{VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
      VkSemaphore signal_semaphores[]{
	BluCat::core.vk_swapchain->render_finished_semaphores[
	  BluCat::core.vk_swapchain->current_frame]};

      VkSubmitInfo submit_info{};
      submit_info.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
      submit_info.pNext = nullptr;
      submit_info.waitSemaphoreCount = 1;
      submit_info.pWaitSemaphores = wait_semaphores;
      submit_info.pWaitDstStageMask = wait_stages;
      submit_info.commandBufferCount = 1;
      submit_info.pCommandBuffers = &draw_command_buffer;
      submit_info.signalSemaphoreCount = 1;
      submit_info.pSignalSemaphores = signal_semaphores;

      if(vkQueueSubmit(
	   queue.queue, 1, &submit_info, BluCat::core.vk_swapchain->in_flight_fences[
	     BluCat::core.vk_swapchain->current_frame]) != VK_SUCCESS)
	throw std::runtime_error{"Failed to submit draw command buffer."};

      VkSwapchainKHR swap_chains[]{BluCat::core.vk_swapchain->swapchain};

      VkPresentInfoKHR present_info{};
      present_info.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
      present_info.pNext = nullptr;
      present_info.waitSemaphoreCount = 1;
      present_info.pWaitSemaphores = signal_semaphores;
      present_info.swapchainCount = 1;
      present_info.pSwapchains = swap_chains;
      present_info.pImageIndices = &image_index;
      present_info.pResults = nullptr;

      vkQueuePresentKHR(queue.queue, &present_info);
    }

    // Prepare for the next frame.
    {
      this->skeletal_models_to_draw[next_frame].clear();
      this->static_models_to_draw[next_frame].clear();
      this->sprites_3d_to_draw[next_frame].clear();
      BluCat::core.vk_swapchain->current_frame = next_frame;
    }
  }
  else
  {
    // Clear images for the current frame because we are skipping this frame.
    this->skeletal_models_to_draw[
			BluCat::core.vk_swapchain->current_frame].clear();
    this->static_models_to_draw[
			BluCat::core.vk_swapchain->current_frame].clear();
    this->sprites_3d_to_draw[
			BluCat::core.vk_swapchain->current_frame].clear();
    for(auto &view: this->views_2d)
      view->sprites_to_draw[BluCat::core.vk_swapchain->current_frame].clear();
    for(auto &view: this->views_3d)
      view->sprites_to_draw[BluCat::core.vk_swapchain->current_frame].clear();
  }
}

}