#include <stdafx.h>
#include "Device.h"
#include "Context.h"


namespace fpr
{
void fpr::Device::CreateLogicalDevice()
{
  std::vector<vk::DeviceQueueCreateInfo> queue_create_infos;

  const std::set<uint32_t> queue_family_indices = { m_graphics_queue_index,
                                                    m_presentation_queue_index,
                                                    m_compute_queue_index };

  //Create queues based on the previously acquired queue indices.
  for(const auto& queue_family_index : queue_family_indices)
  {
    vk::DeviceQueueCreateInfo queue_create_info{};
    queue_create_info.queueFamilyIndex = queue_family_index;
    queue_create_info.queueCount       = 1;

    float priority                     = 1.0f;
    queue_create_info.pQueuePriorities = &priority;
    queue_create_infos.push_back(queue_create_info);
  }

  vk::DeviceCreateInfo device_create_info{};
  device_create_info.queueCreateInfoCount    = (uint32_t)queue_create_infos.size();
  device_create_info.pQueueCreateInfos       = queue_create_infos.data();
  device_create_info.enabledExtensionCount   = (uint32_t)required_extensions.size();
  device_create_info.ppEnabledExtensionNames = required_extensions.data();

  vk::PhysicalDeviceFeatures enabled_features;
  enabled_features.samplerAnisotropy = true;
  device_create_info.setPEnabledFeatures(&enabled_features);

  auto [result_signal, device] = m_physical_device.createDeviceUnique(device_create_info);

  assert(("Logical device creation failed!", result_signal == vk::Result::eSuccess));
  m_logical_device = std::move(device);

  m_graphics_queue     = m_logical_device->getQueue(m_graphics_queue_index, 0);
  m_presentation_queue = m_logical_device->getQueue(m_presentation_queue_index, 0);
  m_compute_queue      = m_logical_device->getQueue(m_compute_queue_index, 0);
}

bool Device::IsDeviceSuitable(vk::PhysicalDevice device, vk::SurfaceKHR surface) const FPR_NOEXCEPT
{
  //Check the device against the requested features and properties.
  vk::PhysicalDeviceProperties device_properties = device.getProperties();
  vk::PhysicalDeviceFeatures   device_features   = device.getFeatures();

  std::vector<vk::QueueFamilyProperties> queue_fam_props = device.getQueueFamilyProperties();

  bool extensions_supported = IsRequiredExtensionsSupported(device);

  bool surface_supported = DeviceSupportsSurface(device, surface);
  return extensions_supported && surface_supported;
}

bool Device::DeviceSupportsSurface(vk::PhysicalDevice device, vk::SurfaceKHR surface) const FPR_NOEXCEPT
{
  auto [capabilities_result, surface_capabilities] = device.getSurfaceCapabilitiesKHR(surface);
  auto [formats_result, formats]                   = device.getSurfaceFormatsKHR(surface);
  auto [modes_result, present_modes]               = device.getSurfacePresentModesKHR(surface);
  return !present_modes.empty() && !formats.empty();
}

vk::SurfaceFormatKHR Device::QueryBestSurfaceFormat(const std::vector<vk::SurfaceFormatKHR>& formats)
{
  //Check for best supported surface format.
  const auto wanted_color_space = vk::ColorSpaceKHR::eSrgbNonlinear;

  for(const auto& available_format : formats)
  {
    for(const auto& wanted_format : wanted_formats)
    {
      if(wanted_format == available_format.format && wanted_color_space == available_format.colorSpace)
      {
        return available_format;
      }
    }
  }

  if(formats.empty())
    return vk::SurfaceFormatKHR{};
  else
    return formats[0]; //Assumed to be ordered in terms of preference.
}

bool Device::IsRequiredExtensionsSupported(vk::PhysicalDevice device) const
{
  auto [enum_result, available_extensions] = device.enumerateDeviceExtensionProperties();
  bool has_extension = false;
  for(const auto& device_extension : required_extensions)
  {
    for(const auto& extension : available_extensions)
    {
      if(strcmp(device_extension, extension.extensionName) == 0)
      {
        has_extension = true;
        break;
      }
    }
    if(!has_extension)
      return false;
  }
  return true;
}

void Device::FindQueueFamilyIndices(vk::SurfaceKHR surface)
{
  std::vector<vk::QueueFamilyProperties> queue_family_props = m_physical_device.getQueueFamilyProperties();
  uint32_t                               queue_index        = 0;

  for(const auto& queue_family : queue_family_props)
  {
    if(queue_family.queueCount > 0 && (queue_family.queueFlags & vk::QueueFlagBits::eGraphics)) 
    {
      m_graphics_queue_index = queue_index;
      if(queue_family.queueFlags & vk::QueueFlagBits::eCompute)
      {
        m_compute_queue_index = queue_index;
      }
    }


    auto [support_result, presentation_support] = m_physical_device.getSurfaceSupportKHR(queue_index, surface);

    if(queue_family.queueCount > 0 && presentation_support)
    {
      m_presentation_queue_index = queue_index;
    }


    ++queue_index;
  }
}

Device::Device(vk::SurfaceKHR surface, vk::Instance& instance)
{
  GetBestDevice(surface, instance);
  FindQueueFamilyIndices(surface);
  CreateLogicalDevice();
}

uint32_t Device::GetGraphicsQueueIndex() const FPR_NOEXCEPT
{
  return m_graphics_queue_index;
}

uint32_t Device::GetPresentQueueIndex() const FPR_NOEXCEPT
{
  return m_presentation_queue_index;
}

uint32_t Device::GetComputeQueueIndex() const FPR_NOEXCEPT
{
  return m_compute_queue_index;
}

const vk::Queue& Device::GetGraphicsQueue() const FPR_NOEXCEPT
{
  return m_graphics_queue;
}

const vk::Queue& Device::GetPresentQueue() const FPR_NOEXCEPT
{
  return m_presentation_queue;
}

const vk::Queue& Device::GetComputeQueue() const FPR_NOEXCEPT
{
  return m_compute_queue;
}

vk::PhysicalDevice& Device::GetPhysicalDeviceHandle() FPR_NOEXCEPT
{
  return m_physical_device;
}

vk::Device& Device::GetLogicalDeviceHandle() FPR_NOEXCEPT
{
  return m_logical_device.get();
}


void Device::GetBestDevice(vk::SurfaceKHR surface, vk::Instance& instance)
{
  auto [enum_result, devices] = instance.enumeratePhysicalDevices();

  assert(("Failed to enumerate physical devices!", enum_result == vk::Result::eSuccess));

  for(auto& device : devices)
  {
    //if(device.getProperties().deviceType != vk::PhysicalDeviceType::eDiscreteGpu)
    //  continue;
    if(IsDeviceSuitable(device, surface))
    {
      m_physical_device = device;
      return;
    }
  }

  std::cout << "No valid device found. Exiting.\n";
  exit(EXIT_FAILURE);

}
} // namespace fpr