#include <set>

#include "precomp.h"
#include "vulkan_engine/vulkan/graphics.h"

namespace veng {

Graphics::QueueFamilyIndices Graphics::FindQueueFamilies(
    VkPhysicalDevice device) {
  std::uint32_t queue_familiy_count = 0;
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queue_familiy_count,
                                           nullptr);
  std::vector<VkQueueFamilyProperties> families(queue_familiy_count);
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queue_familiy_count,
                                           families.data());

  auto graphics_family_it =
      std::find_if(families.begin(), families.end(),
                   [](const VkQueueFamilyProperties& props) {
                     return props.queueFlags &
                            (VK_QUEUE_GRAPHICS_BIT | VK_QUEUE_TRANSFER_BIT);
                   });

  QueueFamilyIndices result;
  result.graphics_family = graphics_family_it - families.begin();

  for (std::uint32_t i = 0; i < families.size(); i++) {
    VkBool32 has_presentation_support = false;
    vkGetPhysicalDeviceSurfaceSupportKHR(device, i, surface_,
                                         &has_presentation_support);
    if (has_presentation_support) {
      result.presentation_family = i;
      break;
    }
  }

  return result;
}

Graphics::SwapChainProperties Graphics::GetSwapChainProperties(
    VkPhysicalDevice device) {
  SwapChainProperties properties;

  vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, surface_,
                                            &properties.capabilities);

  std::uint32_t format_count;
  vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface_, &format_count,
                                       nullptr);
  properties.formats.resize(format_count);
  vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface_, &format_count,
                                       properties.formats.data());

  std::uint32_t modes_count;
  vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface_, &modes_count,
                                            nullptr);
  properties.present_modes.resize(modes_count);
  vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface_, &modes_count,
                                            properties.present_modes.data());

  return properties;
}

std::vector<VkExtensionProperties> Graphics::GetDeviceAvailableExtensions(
    VkPhysicalDevice device) {
  std::uint32_t available_extentions_count;
  vkEnumerateDeviceExtensionProperties(device, nullptr,
                                       &available_extentions_count, nullptr);
  std::vector<VkExtensionProperties> available_extentions(
      available_extentions_count);
  vkEnumerateDeviceExtensionProperties(device, nullptr,
                                       &available_extentions_count,
                                       available_extentions.data());
  return available_extentions;
}

bool Graphics::AreAllDeviceExtensionsSupported(VkPhysicalDevice device) {
  std::vector<VkExtensionProperties> available_extentions =
      GetDeviceAvailableExtensions(device);
  return std::all_of(
      required_device_extentions_.begin(), required_device_extentions_.end(),
      std::bind_front(IsExtensionSupported, available_extentions));
}

bool Graphics::IsDeviceSuitable(VkPhysicalDevice device) {
  QueueFamilyIndices families = FindQueueFamilies(device);
  return families.IsValid() && AreAllDeviceExtensionsSupported(device) &&
         GetSwapChainProperties(device).IsValid();
}

void Graphics::PickPhysicalDevice() {
  std::vector<VkPhysicalDevice> devices = GetAvailableDevices();

  std::erase_if(
      devices, std::not_fn(std::bind_front(&Graphics::IsDeviceSuitable, this)));

  if (devices.empty()) {
    spdlog::error("No physical devices that match the criteria");
    std::exit(EXIT_FAILURE);
  }

  physical_device_ = devices[0];
}

std::vector<VkPhysicalDevice> Graphics::GetAvailableDevices() {
  std::uint32_t device_count;
  vkEnumeratePhysicalDevices(instance_, &device_count, nullptr);

  if (device_count == 0) return {};

  std::vector<VkPhysicalDevice> devices(device_count);
  vkEnumeratePhysicalDevices(instance_, &device_count, devices.data());

  return devices;
}

void Graphics::CreateLogicalDeviceAndQueues() {
  QueueFamilyIndices picked_device_families =
      FindQueueFamilies(physical_device_);

  if (!picked_device_families.IsValid()) {
    std::exit(EXIT_FAILURE);
  }

  std::set<std::uint32_t> unique_queue_families = {
      picked_device_families.graphics_family.value(),
      picked_device_families.presentation_family.value(),
  };

  std::float_t queue_priority = 1.f;

  std::vector<VkDeviceQueueCreateInfo> queue_create_infos;
  for (std::uint32_t unique_queue_family : unique_queue_families) {
    VkDeviceQueueCreateInfo queue_info = {};
    queue_info.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queue_info.queueFamilyIndex = unique_queue_family;
    queue_info.queueCount = 1;
    queue_info.pQueuePriorities = &queue_priority;
    queue_create_infos.push_back(queue_info);
  }

  VkPhysicalDeviceFeatures required_features = {};
  VkPhysicalDeviceFeatures supported_features;
  vkGetPhysicalDeviceFeatures(physical_device_, &supported_features);

  if (supported_features.depthBounds) {
      required_features.depthBounds = VK_TRUE;
  }
  if (supported_features.depthClamp) {
      required_features.depthClamp = VK_TRUE;
  }

  VkDeviceCreateInfo device_info = {};
  device_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
  device_info.queueCreateInfoCount = queue_create_infos.size();
  device_info.pQueueCreateInfos = queue_create_infos.data();
  device_info.pEnabledFeatures = &required_features;
  device_info.enabledExtensionCount = required_device_extentions_.size();
  device_info.ppEnabledExtensionNames = required_device_extentions_.data();

  VkResult result =
      vkCreateDevice(physical_device_, &device_info, nullptr, &logical_device_);
  if (result != VK_SUCCESS) std::exit(EXIT_FAILURE);

  vkGetDeviceQueue(logical_device_,
                   picked_device_families.graphics_family.value(), 0,
                   &graphics_queue_);
  vkGetDeviceQueue(logical_device_,
                   picked_device_families.presentation_family.value(), 0,
                   &present_queue_);
}

}  // namespace veng