#include "swapchain.hpp"

#include <algorithm>
#include <limits>
#include <sophia/util/logger.hpp>

namespace sophia {

Swapchain::Swapchain(Device& device, vk::Extent2D extent)
    : device{device}, extent{extent} {
  initialize();
}

Swapchain::Swapchain(Device& device, vk::Extent2D extent,
                     std::shared_ptr<Swapchain> previous)
    : device{device}, extent{extent}, oldSwapchain{previous} {
  initialize();
  this->oldSwapchain = nullptr;
}

Swapchain::~Swapchain() {
  this->images.clear();
  this->depthImages.clear();

  this->device.get()->destroySwapchainKHR(this->swapchain);
  // log::trace("destroyed vk::SwapchainKHR");
}

vk::Result Swapchain::acquireNextImage(vk::Semaphore signalSemaphore,
                                       u32& imageIndex) {
  return this->device.get()->acquireNextImageKHR(
      this->swapchain, std::numeric_limits<u64>::max(), signalSemaphore,
      VK_NULL_HANDLE, &imageIndex);
}

vk::Result Swapchain::present(vk::Queue presentQueue, u32 imageIndex,
                              vk::Semaphore waitSemaphore) {
  vk::PresentInfoKHR presentInfo = {};
  presentInfo.waitSemaphoreCount = 1;
  presentInfo.pWaitSemaphores = &waitSemaphore;

  vk::SwapchainKHR swapChains[] = {this->swapchain};
  presentInfo.swapchainCount = 1;
  presentInfo.pSwapchains = swapChains;
  presentInfo.pImageIndices = &imageIndex;

  return presentQueue.presentKHR(&presentInfo);
}

void Swapchain::initialize() {
  setDefaultCreateInfo();
  createSwapchain();
  createImages();
  createDepthImages();
}

void Swapchain::setDefaultCreateInfo() {
  SwapchainSupportDetails swapchainSupport = this->device.getSwapchainSupport();

  vk::SurfaceFormatKHR surfaceFormat =
      chooseSurfaceFormat(swapchainSupport.formats);
  vk::PresentModeKHR presentMode =
      choosePresentMode(swapchainSupport.presentModes);
  this->extent = chooseExtent(swapchainSupport.capabilities);

  u32 imageCount = swapchainSupport.capabilities.minImageCount + 1;
  if (swapchainSupport.capabilities.maxImageCount > 0 &&
      imageCount > swapchainSupport.capabilities.maxImageCount) {
    imageCount = swapchainSupport.capabilities.maxImageCount;
  }

  this->swapchainCreateInfo = {vk::SwapchainCreateFlagsKHR(),
                               this->device.getSurface(),
                               imageCount,
                               surfaceFormat.format,
                               surfaceFormat.colorSpace,
                               this->extent,
                               1,
                               vk::ImageUsageFlagBits::eColorAttachment};

  QueueFamilyIndices indices = this->device.getQueueIndices();
  u32 queueFamilyIndices[] = {indices.graphicsFamily.value(),
                              indices.presentFamily.value()};

  if (indices.graphicsFamily != indices.presentFamily) {
    swapchainCreateInfo.imageSharingMode = vk::SharingMode::eConcurrent;
    swapchainCreateInfo.queueFamilyIndexCount = 2;
    swapchainCreateInfo.pQueueFamilyIndices = queueFamilyIndices;
  } else {
    swapchainCreateInfo.imageSharingMode = vk::SharingMode::eExclusive;
  }

  swapchainCreateInfo.preTransform =
      swapchainSupport.capabilities.currentTransform;
  swapchainCreateInfo.compositeAlpha = vk::CompositeAlphaFlagBitsKHR::eOpaque;
  swapchainCreateInfo.presentMode = presentMode;
  swapchainCreateInfo.clipped = VK_TRUE;

  swapchainCreateInfo.oldSwapchain = vk::SwapchainKHR(nullptr);
  swapchainCreateInfo.oldSwapchain = this->oldSwapchain == nullptr
                                         ? VK_NULL_HANDLE
                                         : this->oldSwapchain->swapchain;

  this->imageFormat = surfaceFormat.format;
}

void Swapchain::createSwapchain() {
  try {
    this->swapchain =
        this->device.get()->createSwapchainKHR(this->swapchainCreateInfo);
    // log::trace("created vk::SwapchainKHR");
  } catch (const vk::SystemError& err) {
    log::fatal("failed to create swapchain. Error: ", err.what());
    throw std::runtime_error("failed to create swapchain");
  }
}

void Swapchain::createImages() {
  std::vector<vk::Image> swapchainImages =
      this->device.get()->getSwapchainImagesKHR(this->swapchain);

  this->images.clear();
  this->images.reserve(swapchainImages.size());

  for (vk::Image image : swapchainImages) {
    this->images.push_back(Image::Builder(this->device)
                               .wrapExisting(image)
                               .setFormat(this->imageFormat)
                               .setExtent(width(), height())
                               .withView()
                               .build());
  }

  // log::trace("created all swapchain vk::ImageView");
}

void Swapchain::createDepthImages() {
  this->depthFormat = findDepthFormat();

  this->depthImages.clear();
  this->depthImages.reserve(imageCount());

  for (size_t i = 0; i < imageCount(); i++) {
    this->depthImages.push_back(Image::Builder(this->device)
                                    .asDepthAttachment()
                                    .setFormat(this->depthFormat)
                                    .setExtent(width(), height())
                                    .withView()
                                    .build());
  }

  // log::trace("created all depth resources");
}

vk::SurfaceFormatKHR Swapchain::chooseSurfaceFormat(
    const std::vector<vk::SurfaceFormatKHR>& availableFormats) {
  if (availableFormats.size() == 1 &&
      availableFormats[0].format == vk::Format::eUndefined) {
    return {vk::Format::eB8G8R8A8Unorm, vk::ColorSpaceKHR::eSrgbNonlinear};
  }

  for (const auto& availableFormat : availableFormats) {
    if (availableFormat.format == vk::Format::eB8G8R8A8Unorm &&
        availableFormat.colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear) {
      return availableFormat;
    }
  }

  log::info("No available surface formats, choosing default.");
  return availableFormats[0];
}

vk::PresentModeKHR Swapchain::choosePresentMode(
    const std::vector<vk::PresentModeKHR> availablePresentModes) {
  vk::PresentModeKHR bestMode = vk::PresentModeKHR::eFifo;

  for (const auto& availablePresentMode : availablePresentModes) {
    if (availablePresentMode == vk::PresentModeKHR::eMailbox) {
      log::info("Selected present mode: MAILBOX (triple buffering)");
      return availablePresentMode;
    } else if (availablePresentMode == vk::PresentModeKHR::eImmediate) {
      bestMode = availablePresentMode;
    }
  }
  log::warning("Present mode MAILBOX not available.");

  if (bestMode == vk::PresentModeKHR::eImmediate) {
    log::info("Selected present mode: IMMEDIATE (no vsync)");
  } else {
    log::info("Selected present mode: FIFO (vsync)");
  }

  return bestMode;
}

vk::Extent2D Swapchain::chooseExtent(
    const vk::SurfaceCapabilitiesKHR& capabilities) {
  if (capabilities.currentExtent.width != std::numeric_limits<u32>::max()) {
    return capabilities.currentExtent;
  } else {
    vk::Extent2D choosenExtent = {};
    choosenExtent.width =
        std::clamp(this->extent.width, capabilities.minImageExtent.width,
                   capabilities.maxImageExtent.width);
    choosenExtent.height =
        std::clamp(this->extent.height, capabilities.minImageExtent.height,
                   capabilities.maxImageExtent.height);
    return choosenExtent;
  }
}

vk::Format Swapchain::findDepthFormat() {
  return this->device.findSupportedFormat(
      {vk::Format::eD32Sfloat, vk::Format::eD32SfloatS8Uint,
       vk::Format::eD24UnormS8Uint},
      vk::ImageTiling::eOptimal,
      vk::FormatFeatureFlagBits::eDepthStencilAttachment);
}

}  // namespace sophia