#include <fmt/format.h>
#include <set>
#include <tuple>

#include "device_manager.hpp"

#include "../structs/swap_chain_support_details.hpp"

fn DeviceManager::pickPhysicalDevice(
  const vk::Instance&             instance,
  const vk::SurfaceKHR&           surface,
  const std::vector<const char*>& requiredExtensions
) -> std::tuple<vk::PhysicalDevice, vk::SampleCountFlagBits, f32, bool> {
  // Get all available physical devices
  std::vector<vk::PhysicalDevice> devices = instance.enumeratePhysicalDevices();

  if (devices.empty())
    throw std::runtime_error("Failed to find GPUs with Vulkan support!");

#ifndef NDEBUG
  fmt::println("Available devices:");
#endif

  vk::PhysicalDevice      physicalDevice;
  vk::SampleCountFlagBits msaaSamples     = vk::SampleCountFlagBits::e1;
  f32                     maxLineWidth    = 1.0F;
  bool                    wideLineSupport = false;

  // For each device,
  for (const vk::PhysicalDevice& device : devices) {
#ifndef NDEBUG
    vk::PhysicalDeviceProperties properties = device.getProperties();
    fmt::println("\t{}", properties.deviceName.data());
#endif

    // Set the first suitable device as the physical device
    if (isDeviceSuitable(device, surface, requiredExtensions)) {
      physicalDevice = device;
      msaaSamples    = getMaxUsableSampleCount(physicalDevice);

      // Get the device properties for line width limits
      vk::PhysicalDeviceProperties deviceProperties = device.getProperties();
      maxLineWidth                                  = deviceProperties.limits.lineWidthRange[1];
      wideLineSupport                               = deviceProperties.limits.lineWidthRange[1] > 1.0F;

#ifndef NDEBUG
      fmt::println("Maximum supported line width: {}", maxLineWidth);
      fmt::println("Wide lines supported: {}", wideLineSupport ? "yes" : "no");
#endif
      break;
    }
  }

  // If no suitable device was found, throw an error
  if (!physicalDevice)
    throw std::runtime_error("Failed to find a suitable GPU!");

  return std::make_tuple(physicalDevice, msaaSamples, maxLineWidth, wideLineSupport);
}

fn DeviceManager::createLogicalDevice(
  const vk::PhysicalDevice&       physicalDevice,
  const vk::SurfaceKHR&           surface,
  const std::vector<const char*>& requiredExtensions,
  const bool                      enableValidationLayers,
  const std::vector<const char*>& validationLayers
) -> std::tuple<vk::UniqueDevice, vk::Queue, vk::Queue> {
  QueueFamilyIndices indices = findQueueFamilies(physicalDevice, surface);

  // Create a set of unique queue families needed
  std::set<u32> uniqueQueueFamilies = { indices.graphics_family.value(), indices.present_family.value() };

  // Create queue create infos for each unique queue family
  std::vector<vk::DeviceQueueCreateInfo> queueCreateInfos;
  f32                                    queuePriority = 1.0F;

  for (u32 queueFamily : uniqueQueueFamilies) {
    vk::DeviceQueueCreateInfo queueCreateInfo {
      .queueFamilyIndex = queueFamily,
      .queueCount       = 1,
      .pQueuePriorities = &queuePriority,
    };
    queueCreateInfos.push_back(queueCreateInfo);
  }

  // Specify device features
  vk::PhysicalDeviceFeatures deviceFeatures {
    .fillModeNonSolid  = true,
    .wideLines         = true,
    .samplerAnisotropy = true,
  };

  // Create the logical device
  vk::DeviceCreateInfo createInfo {
    .queueCreateInfoCount    = static_cast<u32>(queueCreateInfos.size()),
    .pQueueCreateInfos       = queueCreateInfos.data(),
    .enabledExtensionCount   = static_cast<u32>(requiredExtensions.size()),
    .ppEnabledExtensionNames = requiredExtensions.data(),
    .pEnabledFeatures        = &deviceFeatures,
  };

  // Only enable validation layers if requested
  if (enableValidationLayers) {
    createInfo.enabledLayerCount   = static_cast<u32>(validationLayers.size());
    createInfo.ppEnabledLayerNames = validationLayers.data();
  }

  // Create the logical device
  vk::UniqueDevice device = physicalDevice.createDeviceUnique(createInfo);

  // Get queue handles
  vk::Queue graphicsQueue = device->getQueue(indices.graphics_family.value(), 0);
  vk::Queue presentQueue  = device->getQueue(indices.present_family.value(), 0);

  return std::make_tuple(std::move(device), graphicsQueue, presentQueue);
}

fn DeviceManager::getMaxUsableSampleCount(const vk::PhysicalDevice& physicalDevice
) -> vk::SampleCountFlagBits {
  vk::PhysicalDeviceProperties physicalDeviceProperties = physicalDevice.getProperties();

  vk::SampleCountFlags counts = physicalDeviceProperties.limits.framebufferColorSampleCounts &
                                physicalDeviceProperties.limits.framebufferDepthSampleCounts;

  if (counts & vk::SampleCountFlagBits::e64)
    return vk::SampleCountFlagBits::e64;
  if (counts & vk::SampleCountFlagBits::e32)
    return vk::SampleCountFlagBits::e32;
  if (counts & vk::SampleCountFlagBits::e16)
    return vk::SampleCountFlagBits::e16;
  if (counts & vk::SampleCountFlagBits::e8)
    return vk::SampleCountFlagBits::e8;
  if (counts & vk::SampleCountFlagBits::e4)
    return vk::SampleCountFlagBits::e4;
  if (counts & vk::SampleCountFlagBits::e2)
    return vk::SampleCountFlagBits::e2;

  return vk::SampleCountFlagBits::e1;
}

fn DeviceManager::findQueueFamilies(const vk::PhysicalDevice& device, const vk::SurfaceKHR& surface)
  -> QueueFamilyIndices {
  QueueFamilyIndices indices;

  // Get queue family properties
  std::vector<vk::QueueFamilyProperties> queueFamilies = device.getQueueFamilyProperties();

  // Find queue family with graphics support
  for (size_t i = 0; i < queueFamilies.size(); i++) {
    if (queueFamilies[i].queueFlags & vk::QueueFlagBits::eGraphics)
      indices.graphics_family = static_cast<u32>(i);

    // Check for presentation support
    if (device.getSurfaceSupportKHR(static_cast<u32>(i), surface))
      indices.present_family = static_cast<u32>(i);

    if (indices.isComplete())
      break;
  }

  return indices;
}

fn DeviceManager::isDeviceSuitable(
  const vk::PhysicalDevice&       device,
  const vk::SurfaceKHR&           surface,
  const std::vector<const char*>& requiredExtensions
) -> bool {
  // Check queue family support
  QueueFamilyIndices indices = findQueueFamilies(device, surface);

  // Check extension support
  bool extensionsSupported = checkDeviceExtensionSupport(device, requiredExtensions);

  // Check swap chain support
  bool swapChainAdequate = false;
  if (extensionsSupported) {
    SwapChainSupportDetails swapChainSupport;
    swapChainSupport.capabilities  = device.getSurfaceCapabilitiesKHR(surface);
    swapChainSupport.formats       = device.getSurfaceFormatsKHR(surface);
    swapChainSupport.present_modes = device.getSurfacePresentModesKHR(surface);
    swapChainAdequate = !swapChainSupport.formats.empty() && !swapChainSupport.present_modes.empty();
  }

  // Check device features
  vk::PhysicalDeviceFeatures supportedFeatures = device.getFeatures();

  return indices.isComplete() && extensionsSupported && swapChainAdequate &&
         supportedFeatures.samplerAnisotropy;
}

fn DeviceManager::checkDeviceExtensionSupport(
  const vk::PhysicalDevice&       device,
  const std::vector<const char*>& requiredExtensions
) -> bool {
  // Get available extensions
  std::vector<vk::ExtensionProperties> availableExtensions = device.enumerateDeviceExtensionProperties();

  // Convert available extensions to a set of strings for easier lookup
  std::set<std::string> availableExtensionNames;
  for (const auto& extension : availableExtensions) {
    availableExtensionNames.insert(extension.extensionName);
  }

  // Check if all required extensions are available
  return std::ranges::all_of(requiredExtensions, [&](const char* required) {
    return availableExtensionNames.contains(required);
  });
}