#include "SwapChain.hpp"

#include <stdexcept>
#include <numeric>
#include <algorithm>

void SwapChain::destroySwapChain(VkDevice device) {
  for (const auto &buffer: mSwapChainFrameBuffers) {
    vkDestroyFramebuffer(device, buffer, nullptr);
  }

  for (const auto &imageView: mSwapChainImageViews) {
    vkDestroyImageView(device, imageView, nullptr);
  }

  vkDestroySwapchainKHR(device, mSwapChain, nullptr);
}

void SwapChain::createSwapChain(const CreateInfo &info) {
  SwapChainSupportDetails details = querySwapChainSupportDetails(info.physicalDevice, info.surface);

  VkSurfaceFormatKHR surfaceFormat = pickSwapChainSurfaceFormat(details);
  VkPresentModeKHR presentMode = pickSwapChainPresentMode(details);
  VkExtent2D extent2D = pickSwapChainExtent(details.capabilities, info.sizeX, info.sizeY);

  // +1 so will not have to wait for device to finish frame to query an image to render to
  uint32_t imageCount = details.capabilities.minImageCount + 1;

  if (details.capabilities.maxImageCount > 0 && imageCount > details.capabilities.maxImageCount) {
    imageCount = details.capabilities.maxImageCount;
  }

  VkSwapchainCreateInfoKHR createInfo{
      .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
      .surface = info.surface,
      .minImageCount = imageCount,
      .imageFormat = surfaceFormat.format,
      .imageExtent = extent2D,
      .imageArrayLayers = 1,
      .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
      .preTransform = details.capabilities.currentTransform,
      .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
      .presentMode = presentMode,
      .clipped = VK_TRUE,
      .oldSwapchain = VK_NULL_HANDLE,
  };

  uint32_t queueIndices[] = {info.graphicsQueueFamilyIndex, info.presentationQueueFamilyIndex};

  if (info.graphicsQueueFamilyIndex != info.presentationQueueFamilyIndex) {
    createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
    createInfo.queueFamilyIndexCount = 2;
    createInfo.pQueueFamilyIndices = queueIndices;
  } else {
    createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
  }

  if (vkCreateSwapchainKHR(info.device, &createInfo, nullptr, &mSwapChain) != VK_SUCCESS) {
    throw std::runtime_error("failed to create swap chain");
  }

  mSwapChainExtent = extent2D;
  mSwapChainFormat = surfaceFormat.format;

  uint32_t createdImageCount = 0;
  vkGetSwapchainImagesKHR(info.device, mSwapChain, &createdImageCount, nullptr);
  mSwapChainImages.resize(createdImageCount);
  vkGetSwapchainImagesKHR(info.device, mSwapChain, &createdImageCount, mSwapChainImages.data());
}


void SwapChain::createSwapChainImageViews(VkDevice device) {
  mSwapChainImageViews.resize(mSwapChainImages.size());

  int i = 0;
  for (const auto &image: mSwapChainImages) {

    VkImageViewCreateInfo createInfo{
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .image = image,
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .format = mSwapChainFormat,
        .components = {
            .r = VK_COMPONENT_SWIZZLE_R,
            .g = VK_COMPONENT_SWIZZLE_G,
            .b = VK_COMPONENT_SWIZZLE_B,
            .a = VK_COMPONENT_SWIZZLE_A,
        },
        .subresourceRange = {
            .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
            .baseMipLevel = 0,
            .levelCount = 1,
            .baseArrayLayer = 0,
            .layerCount = 1,
        }
    };

    if (vkCreateImageView(device, &createInfo, nullptr, &mSwapChainImageViews[i]) != VK_SUCCESS) {
      throw std::runtime_error("cannot create image view for the swap chain");
    }

    i++;
  }
}

void SwapChain::createSwapChainFramebuffers(VkDevice device, VkRenderPass renderPass) {
  mSwapChainFrameBuffers.resize(mSwapChainImageViews.size());

  for (size_t i = 0; i < mSwapChainImageViews.size(); i++) {

    VkImageView attachments[] = {
        mSwapChainImageViews[i],
    };

    VkFramebufferCreateInfo framebufferCreateInfo{
        .sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO,
        .renderPass = renderPass,
        .attachmentCount = 1,
        .pAttachments = attachments,
        .width = mSwapChainExtent.width,
        .height = mSwapChainExtent.height,
        .layers = 1,
    };

    if (vkCreateFramebuffer(device, &framebufferCreateInfo, nullptr, &mSwapChainFrameBuffers[i]) != VK_SUCCESS) {
      throw std::runtime_error("failed to create swapchain framebuffers");
    }
  }
}

VkSurfaceFormatKHR SwapChain::pickSwapChainSurfaceFormat(const SwapChainSupportDetails &details) {
  for (const auto &format: details.formats) {
    if (format.format == VK_FORMAT_B8G8R8A8_SRGB && format.colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR) {
      return format;
    }
  }

  return details.formats.front();
}

VkPresentModeKHR SwapChain::pickSwapChainPresentMode(const SwapChainSupportDetails &details) {
  for (const auto &mode: details.presentModes) {
    if (mode == VK_PRESENT_MODE_MAILBOX_KHR) {
      return mode;
    }
  }

  return VK_PRESENT_MODE_FIFO_KHR; // guaranteed exists
}

VkExtent2D SwapChain::pickSwapChainExtent(const VkSurfaceCapabilitiesKHR &capabilities, int sizeX, int sizeY) {
  // if set by vulkan just keep it
  if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max()) {
    return capabilities.currentExtent;
  }

  VkExtent2D out = {(uint32_t) sizeX, (uint32_t) sizeY};

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

  return out;
}

SwapChain::SwapChainSupportDetails
SwapChain::querySwapChainSupportDetails(VkPhysicalDevice device, VkSurfaceKHR surface) {
  SwapChainSupportDetails details;

  vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, surface, &details.capabilities);

  uint32_t formatCount = 0;
  vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, nullptr);
  if (formatCount != 0) {
    details.formats.resize(formatCount);
    vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, details.formats.data());
  }

  uint32_t modeCount = 0;
  vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &modeCount, nullptr);
  if (modeCount != 0) {
    details.presentModes.resize(modeCount);
    vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &modeCount, details.presentModes.data());
  }

  return details;
}

bool SwapChain::checkDeviceSwapChain(VkPhysicalDevice device, VkSurfaceKHR surface) {
  SwapChainSupportDetails details = querySwapChainSupportDetails(device, surface);
  return !(details.presentModes.empty() || details.formats.empty());
}