#include <expected>
#include <fmt/format.h>
#include <set>
#include <span>

#include "device_manager.hpp"

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

struct PhysicalDeviceInfo {
  vk::PhysicalDevice      physical_device;
  vk::SampleCountFlagBits msaa_samples;
  f32                     max_line_width;
  bool                    wide_line_support;
};

struct LogicalDeviceInfo {
  vk::UniqueDevice device;
  vk::Queue        graphics_queue;
  vk::Queue        present_queue;
};

fn DeviceManager::pickPhysicalDevice(
  const vk::Instance&          instance,
  const vk::SurfaceKHR&        surface,
  std::span<const char* const> requiredExtensions
) -> std::expected<PhysicalDeviceInfo, std::string> {
  // Get all available physical devices
  auto devices = instance.enumeratePhysicalDevices();

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

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

  // Find the first suitable device using ranges
  auto suitableDevice = std::ranges::find_if(devices, [&](const auto& device) {
    return isDeviceSuitable(device, surface, requiredExtensions);
  });

  if (suitableDevice == devices.end())
    return std::unexpected("Failed to find a suitable GPU!");

  vk::PhysicalDeviceProperties deviceProperties = suitableDevice->getProperties();

#ifndef NDEBUG
  fmt::println("\t{}", deviceProperties.deviceName.data());
  fmt::println("Maximum supported line width: {}", deviceProperties.limits.lineWidthRange[1]);
  fmt::println("Wide lines supported: {}", deviceProperties.limits.lineWidthRange[1] > 1.0F ? "yes" : "no");
#endif

  return PhysicalDeviceInfo {
    .physical_device   = *suitableDevice,
    .msaa_samples      = getMaxUsableSampleCount(*suitableDevice),
    .max_line_width    = deviceProperties.limits.lineWidthRange[1],
    .wide_line_support = deviceProperties.limits.lineWidthRange[1] > 1.0F,
  };
}

fn DeviceManager::createLogicalDevice(
  const vk::PhysicalDevice&    physicalDevice,
  const vk::SurfaceKHR&        surface,
  std::span<const char* const> requiredExtensions,
  bool                         enableValidationLayers,
  std::span<const char* const> validationLayers
) -> std::expected<LogicalDeviceInfo, std::string> {
  try {
    QueueFamilyIndices indices = findQueueFamilies(physicalDevice, surface);

    std::set<u32> uniqueQueueFamilies = { indices.graphics_family.value(), indices.present_family.value() };

    std::vector<vk::DeviceQueueCreateInfo> queueCreateInfos;
    f32                                    queuePriority = 1.0F;

    // Use ranges to transform unique queue families into create infos
    std::ranges::transform(uniqueQueueFamilies, std::back_inserter(queueCreateInfos), [&](u32 queueFamily) {
      return vk::DeviceQueueCreateInfo {
        .queueFamilyIndex = queueFamily,
        .queueCount       = 1,
        .pQueuePriorities = &queuePriority,
      };
    });

    vk::PhysicalDeviceFeatures deviceFeatures {
      .fillModeNonSolid  = vk::True,
      .wideLines         = vk::True,
      .samplerAnisotropy = vk::True,
    };

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

    if (enableValidationLayers) {
      createInfo.enabledLayerCount   = static_cast<u32>(validationLayers.size());
      createInfo.ppEnabledLayerNames = validationLayers.data();
    }

    auto device        = physicalDevice.createDeviceUnique(createInfo);
    auto graphicsQueue = device->getQueue(indices.graphics_family.value(), 0);
    auto presentQueue  = device->getQueue(indices.present_family.value(), 0);

    return LogicalDeviceInfo {
      .device         = std::move(device),
      .graphics_queue = graphicsQueue,
      .present_queue  = presentQueue,
    };
  } catch (const vk::SystemError& e) {
    return std::unexpected(fmt::format("Failed to create logical device: {}", e.what()));
  }
}

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,
  std::span<const char* const> 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,
  std::span<const char* const> requiredExtensions
) -> bool {
  std::vector<vk::ExtensionProperties> availableExtensions = device.enumerateDeviceExtensionProperties();

  // Use ranges to check if all required extensions are available
  return std::ranges::all_of(requiredExtensions, [&](const char* required) {
    return std::ranges::any_of(availableExtensions, [&](const vk::ExtensionProperties& available) {
      return strcmp(required, available.extensionName) == 0;
    });
  });
}