#include "vulkan_engine/vulkan/graphics.h"
#include <GLFW/glfw3.h>

#include "precomp.h"

namespace veng {

void Graphics::CreateSurface() {
  VkResult result = glfwCreateWindowSurface(instance_, window->GetHandle(),
                                            nullptr, &surface_);
  if (result != VK_SUCCESS) std::exit(EXIT_FAILURE);
}

bool IsRgbaTypeFormat(const VkSurfaceFormatKHR& format_properties) {
  return format_properties.format == VK_FORMAT_R8G8B8A8_SRGB ||
         format_properties.format == VK_FORMAT_B8G8R8A8_SRGB;
}

bool IsSrgbColorSpace(const VkSurfaceFormatKHR& format_properties) {
  return format_properties.colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR;
}

bool IsCorrectFormat(const VkSurfaceFormatKHR& format_properties) {
  return IsSrgbColorSpace(format_properties) &&
         IsRgbaTypeFormat(format_properties);
}

VkSurfaceFormatKHR Graphics::ChooseSwapSurfaceFormat(
    gsl::span<VkSurfaceFormatKHR> formats) {
  if (formats.size() == 1 && formats[0].format == VK_FORMAT_UNDEFINED) {
    return {VkFormat::VK_FORMAT_R8G8B8A8_SRGB,
            VkColorSpaceKHR::VK_COLORSPACE_SRGB_NONLINEAR_KHR};
  }

  auto it = std::find_if(formats.begin(), formats.end(), IsCorrectFormat);
  if (it != formats.end()) return *it;

  for (const VkSurfaceFormatKHR& format : formats) {
    if (format.format == VK_FORMAT_R8G8B8A8_SRGB &&
        format.colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR) {
      return format;
    }
  }

  return formats[0];
}

VkPresentModeKHR Graphics::ChooseSwapPresentMode(
    gsl::span<VkPresentModeKHR> present_modes) {
  constexpr std::array<VkPresentModeKHR, 3> preferred_modes = {
      VK_PRESENT_MODE_MAILBOX_KHR, VK_PRESENT_MODE_IMMEDIATE_KHR,
      VK_PRESENT_MODE_FIFO_KHR};

  for (const auto& preferred : preferred_modes)
    if (std::find(present_modes.begin(), present_modes.end(), preferred) !=
        present_modes.end())
      return preferred;

  return VK_PRESENT_MODE_FIFO_KHR;
}

VkExtent2D Graphics::ChooseSwapExtent(
    const VkSurfaceCapabilitiesKHR& capabilities) {
  constexpr std::uint32_t kInvalidSize =
      std::numeric_limits<std::uint32_t>::max();

  if (capabilities.currentExtent.width != kInvalidSize) {
    return capabilities.currentExtent;
  } else {
    glm::ivec2 size = window->GetFramebufferSize();
    VkExtent2D actual_extent = {static_cast<std::uint32_t>(size.x),
                                static_cast<std::uint32_t>(size.y)};

    actual_extent.width =
        std::clamp(actual_extent.width, capabilities.minImageExtent.width,
                   capabilities.maxImageExtent.width);
    actual_extent.height =
        std::clamp(actual_extent.height, capabilities.minImageExtent.height,
                   capabilities.maxImageExtent.height);

    return actual_extent;
  }
}

std::uint32_t Graphics::ChooseSwapImageCount(
    const VkSurfaceCapabilitiesKHR& capabilities) {
  std::uint32_t image_count = capabilities.minImageCount + 1;
  if (capabilities.maxImageCount > 0 &&
      capabilities.maxImageCount < image_count)
    image_count = capabilities.maxImageCount;

  return image_count;
}

void Graphics::CreateSwapChain() {
  SwapChainProperties properties = GetSwapChainProperties(physical_device_);

  surface_format_ = ChooseSwapSurfaceFormat(properties.formats);
  present_mode_ = ChooseSwapPresentMode(properties.present_modes);
  extent_ = ChooseSwapExtent(properties.capabilities);

  std::uint32_t image_count = ChooseSwapImageCount(properties.capabilities);

  VkSwapchainCreateInfoKHR info = {};
  info.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
  info.surface = surface_;
  info.minImageCount = image_count;
  info.imageFormat = surface_format_.format;
  info.imageColorSpace = surface_format_.colorSpace;
  info.imageExtent = extent_;
  info.imageArrayLayers = 1;
  info.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
  info.presentMode = present_mode_;
  info.preTransform = properties.capabilities.currentTransform;
  info.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
  info.clipped = VK_TRUE;
  info.oldSwapchain = VK_NULL_HANDLE;

  QueueFamilyIndices indices = FindQueueFamilies(physical_device_);

  if (indices.graphics_family != indices.presentation_family) {
    std::array<std::uint32_t, 2> family_indices = {
        indices.graphics_family.value(), indices.presentation_family.value()};
    info.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
    info.queueFamilyIndexCount = family_indices.size();
    info.pQueueFamilyIndices = family_indices.data();
  } else {
    info.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
  }

  VkResult result =
      vkCreateSwapchainKHR(logical_device_, &info, nullptr, &swap_chain_);
  if (result != VK_SUCCESS) std::exit(EXIT_FAILURE);

  std::uint32_t actual_image_count;
  vkGetSwapchainImagesKHR(logical_device_, swap_chain_, &actual_image_count,
                          nullptr);
  swap_chain_images_.resize(actual_image_count);
  vkGetSwapchainImagesKHR(logical_device_, swap_chain_, &actual_image_count,
                          swap_chain_images_.data());
}

VkImageView Graphics::CreateImageView(VkImage image, VkFormat format, VkImageAspectFlags aspect_flag) {
  VkImageViewCreateInfo info = {};
  info.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
  info.image = image;
  info.viewType = VK_IMAGE_VIEW_TYPE_2D;
  info.format = format;
  info.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
  info.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
  info.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
  info.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;
  info.subresourceRange.aspectMask = aspect_flag;
  info.subresourceRange.baseMipLevel = 0;
  info.subresourceRange.levelCount = 1;
  info.subresourceRange.baseArrayLayer = 0;
  info.subresourceRange.layerCount = 1;

  VkImageView view;
  VkResult result = vkCreateImageView(logical_device_, &info, nullptr, &view);
  if (result != VK_SUCCESS) std::exit(EXIT_FAILURE);
  return view;
}

void Graphics::CreateImageViews() {
  swap_chain_image_views_.resize(swap_chain_images_.size());

  auto image_view_it = swap_chain_image_views_.begin();
  for (VkImage image : swap_chain_images_) {
    *image_view_it = CreateImageView(image, surface_format_.format, VK_IMAGE_ASPECT_COLOR_BIT);
    
    std::advance(image_view_it, 1);
  }
}

}  // namespace veng