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

#define VK_ENABLE_BETA_EXTENSIONS
#define VULKAN_HPP_NO_CONSTRUCTORS
#include <vulkan/vulkan.hpp>

#include "util/types.h"
#include "vkfw.hpp"

namespace vk {
  using DebugUtilsMessengerUnique = vk::UniqueHandle<vk::DebugUtilsMessengerEXT, vk::DispatchLoaderDynamic>;
}

constexpr i32 WIDTH  = 800;
constexpr i32 HEIGHT = 600;

constexpr std::array<const char*, 1> validationLayers = { "VK_LAYER_KHRONOS_validation" };

#ifdef __APPLE__
constexpr std::array<const char*, 2> deviceExtensions = { vk::KHRSwapchainExtensionName,
                                                          vk::KHRPortabilitySubsetExtensionName };
#else
constexpr std::array<const char*, 1> deviceExtensions = { vk::KHRSwapchainExtensionName };
#endif

#ifdef NDEBUG
constexpr bool enableValidationLayers = false;
#else
constexpr bool enableValidationLayers = true;
#endif

class VulkanApp {
 public:
  fn run() -> void {
    initWindow();
    initVulkan();
    mainLoop();
  }

 private:
  vkfw::UniqueInstance mGLFWInstance;
  vkfw::UniqueWindow   mWindow;

  vk::UniqueInstance mInstance;

  vk::DebugUtilsMessengerUnique mDebugMessenger;
  vk::DispatchLoaderDynamic     mLoader;

  vk::PhysicalDevice mPhysicalDevice;
  vk::UniqueDevice   mDevice;

  vk::Queue mGraphicsQueue;
  vk::Queue mPresentQueue;

  vk::UniqueSurfaceKHR mSurface;

  struct QueueFamilyIndices {
    std::optional<u32> graphics_family;
    std::optional<u32> present_family;

    fn isComplete() -> bool { return graphics_family.has_value() && present_family.has_value(); }
  };

  struct SwapChainSupportDetails {
    vk::SurfaceCapabilitiesKHR        capabilities;
    std::vector<vk::SurfaceFormatKHR> formats;
    std::vector<vk::PresentModeKHR>   present_modes;
  };

  fn initWindow() -> void {
    mGLFWInstance = vkfw::initUnique();

    vkfw::WindowHints hints;

    hints.clientAPI = vkfw::ClientAPI::eNone;
    hints.resizable = vkfw::eFalse;

    mWindow = vkfw::createWindowUnique(WIDTH, HEIGHT, "Vulkan", hints);
  }

  fn initVulkan() -> void {
    createInstance();
    setupDebugMessenger();
    createSurface();
    pickPhysicalDevice();
    createLogicalDevice();
  }

  fn mainLoop() -> void {
    while (!mWindow->shouldClose()) { vkfw::waitEvents(); }
  }

  fn createInstance() -> void {
    if (enableValidationLayers && !checkValidationLayerSupport())
      throw std::runtime_error("Validation layers requested, but not available!");

    vk::ApplicationInfo appInfo { .pApplicationName   = "Hello Triangle",
                                  .applicationVersion = 1,
                                  .pEngineName        = "No Engine",
                                  .engineVersion      = 1,
                                  .apiVersion         = vk::ApiVersion10 };

    // Retrieve extensions using custom function
    std::vector<const char*> extensions = getRequiredExtensions();

#ifdef __APPLE__
    // Enable the portability extension and set flags
    extensions.emplace_back(vk::KHRPortabilityEnumerationExtensionName);
    // Technically deprecated but vulkan complains if I don't include it for macOS
    // So instead of using the vk::KHRPortabilitySubsetExtensionName, I just use
    // the direct string.
    extensions.emplace_back("VK_KHR_get_physical_device_properties2");
#endif

    vk::InstanceCreateInfo createInfo {
      .flags                   = vk::InstanceCreateFlagBits::eEnumeratePortabilityKHR,
      .pApplicationInfo        = &appInfo,
      .enabledLayerCount       = enableValidationLayers ? static_cast<u32>(validationLayers.size()) : 0,
      .ppEnabledLayerNames     = enableValidationLayers ? validationLayers.data() : nullptr,
      .enabledExtensionCount   = static_cast<u32>(extensions.size()),
      .ppEnabledExtensionNames = extensions.data()
    };

#ifndef NDEBUG
    fmt::println("Available extensions:");

    for (const char* extension : extensions) fmt::println("\t{}", extension);
#endif

    try {
      mInstance = vk::createInstanceUnique(createInfo);
      mLoader   = vk::DispatchLoaderDynamic(mInstance.get(), vkGetInstanceProcAddr);
    } catch (const vk::SystemError& err) {
      throw std::runtime_error("Failed to create instance: " + std::string(err.what()));
    }
  }

  fn setupDebugMessenger() -> void {
    if (!enableValidationLayers)
      return;

    vk::DebugUtilsMessengerCreateInfoEXT messengerCreateInfo {
      .messageSeverity = vk::DebugUtilsMessageSeverityFlagBitsEXT::eVerbose |
                         vk::DebugUtilsMessageSeverityFlagBitsEXT::eWarning |
                         vk::DebugUtilsMessageSeverityFlagBitsEXT::eError,
      .messageType = vk::DebugUtilsMessageTypeFlagBitsEXT::eGeneral |
                     vk::DebugUtilsMessageTypeFlagBitsEXT::eValidation |
                     vk::DebugUtilsMessageTypeFlagBitsEXT::ePerformance,
      .pfnUserCallback = debugCallback,
    };

    mDebugMessenger = mInstance->createDebugUtilsMessengerEXTUnique(messengerCreateInfo, nullptr, mLoader);
  }

  fn createSurface() -> void { mSurface = vkfw::createWindowSurfaceUnique(mInstance.get(), mWindow.get()); }

  fn pickPhysicalDevice() -> void {
    std::vector<vk::PhysicalDevice> devices = mInstance->enumeratePhysicalDevices();

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

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

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

      if (isDeviceSuitable(device)) {
        mPhysicalDevice = device;
        break;
      }
    }

    if (!mPhysicalDevice)
      throw std::runtime_error("Failed to find a suitable GPU!");
  }

  fn createLogicalDevice() -> void {
    QueueFamilyIndices indices = findQueueFamilies(mPhysicalDevice);

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

    f32 queuePriority = 1.0F;

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

      queueCreateInfos.emplace_back(queueCreateInfo);
    }

    vk::PhysicalDeviceFeatures deviceFeatures;

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

    mDevice = mPhysicalDevice.createDeviceUnique(createInfo);

    mGraphicsQueue = mDevice->getQueue(indices.graphics_family.value(), 0);
    mPresentQueue  = mDevice->getQueue(indices.present_family.value(), 0);
  }

  fn isDeviceSuitable(vk::PhysicalDevice device) -> bool {
    QueueFamilyIndices indices = findQueueFamilies(device);

    bool extensionsSupported = checkDeviceExtensionSupport(device);

    bool swapChainAdequate = false;

    if (extensionsSupported) {
      SwapChainSupportDetails swapChainSupport = querySwapChainSupport(device);
      swapChainAdequate = !swapChainSupport.formats.empty() && !swapChainSupport.present_modes.empty();
    }

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

  static fn checkDeviceExtensionSupport(vk::PhysicalDevice device) -> bool {
    std::vector<vk::ExtensionProperties> availableExtensions = device.enumerateDeviceExtensionProperties();

    std::set<string> requiredExtensions(deviceExtensions.begin(), deviceExtensions.end());

    for (const auto& extension : availableExtensions) requiredExtensions.erase(extension.extensionName);

    return requiredExtensions.empty();
  }

  fn findQueueFamilies(vk::PhysicalDevice device) -> QueueFamilyIndices {
    QueueFamilyIndices indices;

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

    for (u32 i = 0; i < queueFamilies.size(); i++) {
      if (queueFamilies[i].queueFlags & vk::QueueFlagBits::eGraphics)
        indices.graphics_family = i;

      vk::Bool32 presentSupport = device.getSurfaceSupportKHR(i, mSurface.get());

      if (presentSupport)
        indices.present_family = i;

      if (indices.isComplete())
        break;
    }

    return indices;
  }

  fn querySwapChainSupport(vk::PhysicalDevice device) -> SwapChainSupportDetails {
    SwapChainSupportDetails details;

    details.capabilities  = device.getSurfaceCapabilitiesKHR(mSurface.get());
    details.formats       = device.getSurfaceFormatsKHR(mSurface.get());
    details.present_modes = device.getSurfacePresentModesKHR(mSurface.get());

    return details;
  }

  static fn getRequiredExtensions() -> std::vector<const char*> {
    std::span<const char*> extensionsSpan = vkfw::getRequiredInstanceExtensions();

    std::vector extensions(extensionsSpan.begin(), extensionsSpan.end());

    if (enableValidationLayers)
      extensions.emplace_back(vk::EXTDebugUtilsExtensionName);

    return extensions;
  }

  static fn checkValidationLayerSupport() -> bool {
    std::vector<vk::LayerProperties> availableLayers = vk::enumerateInstanceLayerProperties();

    for (const char* layerName : validationLayers) {
      bool layerFound = false;

      for (const auto& layerProperties : availableLayers)
        if (strcmp(layerName, layerProperties.layerName) == 0) {
          layerFound = true;
          break;
        }

      if (!layerFound)
        return false;
    }

    return true;
  }

  static VKAPI_ATTR fn VKAPI_CALL debugCallback(
    VkDebugUtilsMessageSeverityFlagBitsEXT /*messageSeverity*/,
    VkDebugUtilsMessageTypeFlagsEXT /*messageType*/,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
    void* /*pUserData*/
  ) -> vk::Bool32 {
    fmt::println("Validation layer: {}", pCallbackData->pMessage);

    return vk::False;
  }
};

fn main() -> i32 {
  VulkanApp app;

  try {
    app.run();
  } catch (const std::exception& e) {
    std::cerr << e.what() << std::endl;
    return EXIT_FAILURE;
  }

  return EXIT_SUCCESS;
}