#include "cve_renderer.hpp"

#include <stdexcept>
#include <array>

namespace cve {
  CveRenderer::CveRenderer(CveWindow& window, CveDevice& device) : cveWindow{window}, cveDevice{device}{
    recreateSwapChain();
    createCommandBuffers();
  }

  CveRenderer::~CveRenderer() {
    freeCommandBuffers();
  }

  void CveRenderer::recreateSwapChain() {
    auto extent = cveWindow.getExtent();
    while (extent.width == 0 || extent.height == 0) {
      extent = cveWindow.getExtent();
      glfwWaitEvents();
    }

    vkDeviceWaitIdle(cveDevice.device());

    if (cveSwapChain == nullptr) {
      cveSwapChain = std::make_unique<CveSwapChain>(cveDevice, extent);
    } else {
      cveSwapChain = std::make_unique<CveSwapChain>(cveDevice, extent, std::move(cveSwapChain));
      if (cveSwapChain->imageCount() != commandBuffers.size()) {
        freeCommandBuffers();
        createCommandBuffers();
      }
    }

    // Come back
  }

  void CveRenderer::createCommandBuffers() {
    commandBuffers.resize(cveSwapChain->imageCount());

    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandPool = cveDevice.getCommandPool();
    allocInfo.commandBufferCount = static_cast<uint32_t>(commandBuffers.size());

    if (vkAllocateCommandBuffers(cveDevice.device(), &allocInfo, commandBuffers.data()) !=
        VK_SUCCESS) {
      throw std::runtime_error("Failed to allocate command buffers");
    }
  }

  void CveRenderer::freeCommandBuffers() {
    vkFreeCommandBuffers(cveDevice.device(), cveDevice.getCommandPool(), 
      static_cast<uint32_t>(commandBuffers.size()), commandBuffers.data());
    commandBuffers.clear();
  }

  VkCommandBuffer CveRenderer::beginFrame() {
    assert(!isFrameStarted && "Can't call begin frame if frame is already in progress");

    auto result = cveSwapChain->acquireNextImage(&currentImageIndex);

    if (result == VK_ERROR_OUT_OF_DATE_KHR) {
      recreateSwapChain();
      return nullptr;
    }
    if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
      throw std::runtime_error("Failed to acquire swap chain image!");
    }

    isFrameStarted = true;

    auto commandBuffer = getCurrentCommandBuffer();

    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

    if (vkBeginCommandBuffer(commandBuffer, &beginInfo) != VK_SUCCESS) {
      throw std::runtime_error("Failed to begin command buffer recording");
    }

    return commandBuffer;
  }
  void CveRenderer::endFrame() {
    assert(isFrameStarted && "Can't end a frame that hasn't started");
    auto commandBuffer = getCurrentCommandBuffer();

    if (vkEndCommandBuffer(commandBuffer) != VK_SUCCESS) {
      throw std::runtime_error("Failed to record command buffer!");
    }

    auto result = cveSwapChain->submitCommandBuffers(&commandBuffer, &currentImageIndex);
    if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || cveWindow.wasWindowResized()) {
      cveWindow.resetWindowResizedFlag();
      recreateSwapChain();
    } else if (result != VK_SUCCESS) {
      throw std::runtime_error("Failed to present swap chain image!");
    }

    isFrameStarted = false;
  }

  void CveRenderer::beginSwapChainRenderPass(VkCommandBuffer commandBuffer) {
    assert(isFrameStarted && "Can't begin swap chain render pass if frame isn't started");
    assert(commandBuffer == getCurrentCommandBuffer() && 
    "Can't begin render pass on buffer from another frame.");

    VkRenderPassBeginInfo renderPassInfo{};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassInfo.renderPass = cveSwapChain->getRenderPass();
    renderPassInfo.framebuffer = cveSwapChain->getFrameBuffer(currentImageIndex);

    renderPassInfo.renderArea.offset = {0, 0};
    renderPassInfo.renderArea.extent = cveSwapChain->getSwapChainExtent();

    std::array<VkClearValue, 2> clearValues{};
    clearValues[0].color = {0.01f, 0.01f, 0.01f, 1.0f};
    clearValues[1].depthStencil = {1.0f, 0};
    renderPassInfo.clearValueCount = static_cast<uint32_t>(clearValues.size());
    renderPassInfo.pClearValues = clearValues.data();

    vkCmdBeginRenderPass(commandBuffer, &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);

    VkViewport viewport{};
    viewport.x = 0.0f;
    viewport.y = 0.0f;
    viewport.width = static_cast<float>(cveSwapChain->getSwapChainExtent().width);
    viewport.height = static_cast<float>(cveSwapChain->getSwapChainExtent().height);
    viewport.minDepth = 0.0f;
    viewport.maxDepth = 1.0f;
    VkRect2D scissor{{0, 0}, cveSwapChain->getSwapChainExtent()};
    vkCmdSetViewport(commandBuffer, 0, 1, &viewport);
    vkCmdSetScissor(commandBuffer, 0, 1, &scissor);
  }
  void CveRenderer::endSwapChainRenderPass(VkCommandBuffer commandBuffer) {
    assert(isFrameStarted && "Can't end swap chain render pass if frame isn't started");
    assert(commandBuffer == getCurrentCommandBuffer() && 
    "Can't end render pass on buffer from another frame.");

    vkCmdEndRenderPass(commandBuffer);
  }
}