#include "renderer/render_context.hpp"

#include <limits>
#include <sophia/util/assert.hpp>
#include <sophia/util/logger.hpp>
#include <stdexcept>
#include <tracy/Tracy.hpp>

#include "vulkan/access.hpp"

namespace sophia {

RenderContext::RenderContext(Window& window, Device& device)
    : window{window}, device{device} {
  createFrames();
  recreateSwapchain();
}

RenderContext::~RenderContext() { this->device.waitIdle(); }

vk::CommandBuffer RenderContext::beginFrame() {
  ZoneScoped;
  SOPHIA_ASSERT(!this->isFrameStarted,
                "Can't call beginFrame while already in progress");

  FrameData& frame = this->frames[this->currentFrameIndex];

  vk::Fence inFlightFence = frame.inFlight.get();
  vk::Result fenceResult;
  {
    ZoneScopedN("wait for frame fence");
    fenceResult = this->device.get()->waitForFences(
        inFlightFence, vk::True, std::numeric_limits<u64>::max());
  }

  if (fenceResult != vk::Result::eSuccess) {
    log::error("failed to wait for fence: " + vk::to_string(fenceResult));
  }

  vk::Semaphore imageAvailable = frame.imageAvailable.get();
  vk::Result result = this->swapchain->acquireNextImage(
      imageAvailable, this->currentImageIndex);
  if (result == vk::Result::eErrorOutOfDateKHR) {
    recreateSwapchain();
    return nullptr;
  }

  if (result != vk::Result::eSuccess && result != vk::Result::eSuboptimalKHR) {
    log::fatal("failed to acquire swapchain image");
    throw std::runtime_error("failed to acquire swapchain image");
  }

  this->isFrameStarted = true;

  vk::Fence& imageInFlight = this->imagesInFlight.at(this->currentImageIndex);
  if (imageInFlight) {
    vk::Result fenceResult = this->device.get()->waitForFences(
        imageInFlight, vk::True, std::numeric_limits<u64>::max());

    if (fenceResult != vk::Result::eSuccess) {
      log::error("failed to wait for fence: " + vk::to_string(fenceResult));
    }
  }

  imageInFlight = inFlightFence;

  this->device.get()->resetFences(inFlightFence);

  frame.commandBuffer.reset();

  try {
    frame.commandBuffer.begin(vk::CommandBufferBeginInfo{});
  } catch (const vk::SystemError& err) {
    log::fatal("failed to begin recording command buffer");
    throw std::runtime_error("failed to begin recording command buffer");
  }

  return frame.commandBuffer;
}

void RenderContext::endFrame() {
  ZoneScoped;
  SOPHIA_ASSERT(this->isFrameStarted,
                "Can't call endFrame while no frame is in progress");

  FrameData& frame = this->frames[this->currentFrameIndex];
  frame.commandBuffer.end();

  vk::Semaphore signalSemaphore =
      this->renderFinishedSemaphores.at(this->currentImageIndex).get();

  vk::SemaphoreSubmitInfo waitInfo{};
  waitInfo.semaphore = frame.imageAvailable.get();
  waitInfo.stageMask = SWAPCHAIN_ACQUIRE_STAGE;

  // Present has no pipeline stage to wait on, so signal once everything in
  // the submission is done.
  vk::SemaphoreSubmitInfo signalInfo{};
  signalInfo.semaphore = signalSemaphore;
  signalInfo.stageMask = vk::PipelineStageFlagBits2::eAllCommands;

  vk::CommandBufferSubmitInfo commandBufferInfo{};
  commandBufferInfo.commandBuffer = frame.commandBuffer;

  vk::SubmitInfo2 submitInfo{};
  submitInfo.setWaitSemaphoreInfos(waitInfo);
  submitInfo.setCommandBufferInfos(commandBufferInfo);
  submitInfo.setSignalSemaphoreInfos(signalInfo);

  this->device.getGraphicsQueue().submit2(submitInfo, frame.inFlight.get());

  vk::Result result = this->swapchain->present(
      this->device.getPresentQueue(), this->currentImageIndex, signalSemaphore);
  if (result == vk::Result::eErrorOutOfDateKHR ||
      result == vk::Result::eSuboptimalKHR || this->window.wasResized()) {
    this->window.resetResizedFlag();
    recreateSwapchain();
  } else if (result != vk::Result::eSuccess) {
    log::fatal("failed to present swapchain image");
    throw std::runtime_error("failed to present swapchain image");
  }

  this->isFrameStarted = false;
  this->currentFrameIndex =
      (this->currentFrameIndex + 1) % MAX_FRAMES_IN_FLIGHT;
}

void RenderContext::recreateSwapchain() {
  vk::Extent2D extent = this->window.getExtent();
  while (extent.width == 0 || extent.height == 0) {
    extent = this->window.getExtent();
    glfwWaitEvents();
  }

  this->device.waitIdle();

  if (this->swapchain == nullptr) {
    this->swapchain = std::make_unique<Swapchain>(this->device, extent);
  } else {
    std::shared_ptr<Swapchain> oldSwapchain = std::move(this->swapchain);
    this->swapchain =
        std::make_unique<Swapchain>(this->device, extent, oldSwapchain);

    if (!oldSwapchain->compareSwapchainFormats(*this->swapchain)) {
      log::fatal("swapchain image format changed on recreation");
      throw std::runtime_error("swapchain image format changed on recreation");
    }
  }

  this->renderFinishedSemaphores.clear();
  for (size_t i = 0; i < this->swapchain->imageCount(); i++) {
    this->renderFinishedSemaphores.push_back(
        this->device.get()->createSemaphoreUnique({}));
  }

  // Old fences referred to images owned by the now-replaced swapchain;
  // there's nothing valid left to wait on.
  this->imagesInFlight.assign(this->swapchain->imageCount(), vk::Fence{});
}

void RenderContext::createFrames() {
  vk::CommandBufferAllocateInfo allocInfo{this->device.getCommandPool(),
                                          vk::CommandBufferLevel::ePrimary,
                                          MAX_FRAMES_IN_FLIGHT};
  std::vector<vk::CommandBuffer> commandBuffers =
      this->device.get()->allocateCommandBuffers(allocInfo);

  for (u32 i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    this->frames[i].commandBuffer = commandBuffers[i];
    this->frames[i].imageAvailable =
        this->device.get()->createSemaphoreUnique({});
    this->frames[i].inFlight = this->device.get()->createFenceUnique(
        {vk::FenceCreateFlagBits::eSignaled});
  }
}

}  // namespace sophia