#if defined(__INTELLISENSE__) || !defined(USE_CPP20_MODULES)
#include "vulkan/vulkan.hpp"
#include <vulkan/vulkan_raii.hpp>
#include <vulkan/vulkan_core.h>
#else
import vulkan_hpp;
#endif

#define GLFW_INCLUDE_VULKAN
#include <GLFW/glfw3.h>
#include <ios>
#include <iostream>
#include <fstream>
#include <vector>
#include <string>
#include <stdexcept>
#include <cstdlib>
#include <cstdint>
#include <limits>
#include <algorithm>

constexpr uint32_t WIDTH = 800;
constexpr uint32_t HEIGHT = 600;

const std::vector<char const *> validationLayers = {
  "VK_LAYER_KHRONOS_validation"
};

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

static std::vector<char> readFile(const std::string &filename) {
  std::ifstream file(filename, std::ios::ate | std::ios::binary);
  if (!file.is_open()) {
    throw std::runtime_error("failed to open file!");
  }

  std::vector<char> buffer(file.tellg());
  file.seekg(0, std::ios::beg);
  file.read(buffer.data(), static_cast<std::streamsize>(buffer.size()));
  file.close();

  return buffer;
}

class HelloTriangleApplication {
  public:
    void run() {
      initWindow();
      initVulkan();
      mainLoop();
      cleanup();
    }
  private:
    void initWindow() {
      glfwInit();
      // Don't create an OpenGL context
      glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
      glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
      window = glfwCreateWindow(WIDTH, HEIGHT, "Vulkan", nullptr, nullptr);
    }
    void initVulkan() {
      createInstance();
      createSurface();
      pickPhysicalDevice();
      createLogicalDevice();
      createSwapChain();
      createImageViews();
      createGraphicsPipeline();
    }
    void mainLoop() {
      while (!glfwWindowShouldClose(window)) {
        glfwPollEvents();
      }
    }
    void cleanup() {
      glfwDestroyWindow(window);
      glfwTerminate();
    }
    void createInstance() {
      constexpr vk::ApplicationInfo appInfo {
        .pApplicationName   = "Hello Triangle",
        .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
        .pEngineName        = "No Engine",
        .engineVersion      = VK_MAKE_VERSION(1, 0, 0),
        .apiVersion         = vk::ApiVersion14,
      };

      // Get the required layers
      std::vector<char const*> requiredLayers;
      if (enableValidationLayers) {
        requiredLayers.assign(validationLayers.begin(), validationLayers.end());
      }

      // Check if the required layers are supported by the Vulkan implementation.
      auto layerProperties = context.enumerateInstanceLayerProperties();
      if (std::ranges::any_of(requiredLayers, [&layerProperties](auto const& requiredLayer) {
        return std::ranges::none_of(layerProperties,
                                   [requiredLayer](auto const& layerProperty)
                                   { return strcmp(layerProperty.layerName, requiredLayer) == 0; });
      }))
      {
        throw std::runtime_error("One or more required layers are not supported!");
      }

      // Get the required instance extensions from GLFW.
      uint32_t glfwExtensionCount = 0;
      auto glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);

      // Check if the required GLFW extensions are supported by the Vulkan implementation.
      auto extensionProperties = context.enumerateInstanceExtensionProperties();
      for (uint32_t i = 0; i < glfwExtensionCount; ++i)
      {
        if (std::ranges::none_of(extensionProperties,
                                 [glfwExtension = glfwExtensions[i]](auto const& extensionProperty)
                                 { return strcmp(extensionProperty.extensionName, glfwExtension) == 0; }))
        {
          throw std::runtime_error("Required GLFW extension not supported: " + std::string(glfwExtensions[i]));
        }
      }
      vk::InstanceCreateInfo createInfo {
        .pApplicationInfo        = &appInfo,
        .enabledLayerCount       = static_cast<uint32_t>(requiredLayers.size()),
        .ppEnabledLayerNames     = requiredLayers.data(),
        .enabledExtensionCount   = glfwExtensionCount,
        .ppEnabledExtensionNames = glfwExtensions,
      };

      instance = vk::raii::Instance(context, createInfo);
    }
    void createSurface() {
      VkSurfaceKHR _surface;
      if (glfwCreateWindowSurface(*instance, window, nullptr, &_surface) != 0) {
        throw std::runtime_error("failed to create window surface!");
      }

      surface = vk::raii::SurfaceKHR(instance, _surface);
    }
    void pickPhysicalDevice() {
      std::vector<const char*> deviceExtensions = {
        vk::KHRSwapchainExtensionName,
        vk::KHRSpirv14ExtensionName,
        vk::KHRSynchronization2ExtensionName,
        vk::KHRCreateRenderpass2ExtensionName,
      };

      auto devices = instance.enumeratePhysicalDevices();
      if (devices.empty()) {
        throw std::runtime_error("failed to find GPUs with Vulkan support!");
      }

      for (const auto &device : devices) {
        auto deviceProperties = device.getProperties();
        auto deviceFeatures   = device.getFeatures();
        auto queueFamilies    = device.getQueueFamilyProperties();
        auto extensions       = device.enumerateDeviceExtensionProperties();
        bool isSuitable       = deviceProperties.apiVersion >= VK_API_VERSION_1_3;
        bool extensionFound   = true;

        const vk::QueueFamilyProperties *qf = nullptr;
        for (const auto &qfp : queueFamilies) {
          if ((qfp.queueFlags & vk::QueueFlagBits::eGraphics) != static_cast<vk::QueueFlags>(0)) {
            qf = &qfp;
            break;
          }
        }

        isSuitable = isSuitable && (qf != nullptr);

        for (const auto &extension : deviceExtensions) {
          auto extensionIter = std::ranges::find_if(extensions, [extension](auto const & ext) {return strcmp(ext.extensionName, extension) == 0;});
          extensionFound     = extensionFound &&  extensionIter != extensions.end();
        }

        isSuitable = isSuitable && extensionFound;

        if (isSuitable) {
          physicalDevice = device;
          return;
        }

        throw std::runtime_error("failed to find a suitable GPU");
      }
    }
    void createLogicalDevice() {
      std::vector<vk::QueueFamilyProperties> queueFamilyProperties = physicalDevice.getQueueFamilyProperties();
      uint32_t graphicsIndex = findQueueFamilies(physicalDevice);
      float queuePriority    = 0.5f;
      vk::DeviceQueueCreateInfo deviceQueueCreateInfo {
        .queueFamilyIndex = graphicsIndex,
        .queueCount       = 1,
        .pQueuePriorities = &queuePriority,
      };

      // Create a chain of feature structures
      vk::StructureChain<vk::PhysicalDeviceFeatures2,
                         vk::PhysicalDeviceVulkan13Features,
                         vk::PhysicalDeviceVulkan11Features,
                         vk::PhysicalDeviceExtendedDynamicStateFeaturesEXT> featureChain = {
        {},                               // vk::PhysicalDeviceFeatures2 (empty for now)
        {.dynamicRendering = true },      // Enable dynamic rendering from Vulkan 1.3
        {.shaderDrawParameters = true },  // Enable shader draw parameters from Vulkan 1.2
        {.extendedDynamicState = true }   // Enable extended dynamic state from the extension
      };

      std::vector<const char*> deviceExtensions = {
        vk::KHRSwapchainExtensionName,
        vk::KHRSpirv14ExtensionName,
        vk::KHRSynchronization2ExtensionName,
        vk::KHRCreateRenderpass2ExtensionName
      };

      vk::DeviceCreateInfo deviceCreateInfo {
        .pNext                   = &featureChain.get<vk::PhysicalDeviceFeatures2>(),
        .queueCreateInfoCount    = 1,
        .pQueueCreateInfos       = &deviceQueueCreateInfo,
        .enabledExtensionCount   = static_cast<uint32_t>(deviceExtensions.size()),
        .ppEnabledExtensionNames = deviceExtensions.data(),
      };

      device        = vk::raii::Device(physicalDevice, deviceCreateInfo);
      graphicsQueue = vk::raii::Queue(device, graphicsIndex, 0);
    }
    void createSwapChain() {
      auto surfaceCapabilities = physicalDevice.getSurfaceCapabilitiesKHR(surface);
      swapChainSurfaceFormat   = chooseSwapSurfaceFormat(physicalDevice.getSurfaceFormatsKHR(surface));
      swapChainExtent          = chooseSwapExtent(surfaceCapabilities);
      auto minImageCount       = std::max(3u, surfaceCapabilities.minImageCount);
      minImageCount            = (surfaceCapabilities.maxImageCount > 0 &&
                                  minImageCount > surfaceCapabilities.maxImageCount) ?
                                      surfaceCapabilities.maxImageCount              :
                                      minImageCount;

      vk::SwapchainCreateInfoKHR swapChainCreateInfo {
        .flags            = vk::SwapchainCreateFlagsKHR(),
        .surface          = surface,
        .minImageCount    = minImageCount,
        .imageFormat      = swapChainSurfaceFormat.format,
        .imageColorSpace  = swapChainSurfaceFormat.colorSpace,
        .imageExtent      = swapChainExtent,
        .imageArrayLayers = 1,
        .imageUsage       = vk::ImageUsageFlagBits::eColorAttachment,
        .imageSharingMode = vk::SharingMode::eExclusive,
        .preTransform     = surfaceCapabilities.currentTransform,
        .compositeAlpha   = vk::CompositeAlphaFlagBitsKHR::eOpaque,
        .presentMode      = chooseSwapPresentMode(physicalDevice.getSurfacePresentModesKHR(surface)),
        .clipped          = true,
        .oldSwapchain     = nullptr,
      };

      swapChain            = vk::raii::SwapchainKHR(device, swapChainCreateInfo);
      swapChainImages      = swapChain.getImages();
      swapChainImageFormat = swapChainSurfaceFormat.format;
    }
    void createImageViews() {
      swapChainImageViews.clear();

      vk::ImageViewCreateInfo imageViewCreateInfo{
        .viewType = vk::ImageViewType::e2D,
        .format = swapChainImageFormat,
        .subresourceRange = { vk::ImageAspectFlagBits::eColor, 0, 1, 0, 1 }
      };

      for (auto image : swapChainImages) {
        imageViewCreateInfo.image = image;
        swapChainImageViews.emplace_back(vk::raii::ImageView(device, imageViewCreateInfo));
      }
    }
    void createGraphicsPipeline() {
      vk::raii::ShaderModule shaderModule = createShaderModule(readFile("shaders/09_shader_base.spv"));

      vk::PipelineShaderStageCreateInfo vertShaderStageInfo = {
        .stage  = vk::ShaderStageFlagBits::eVertex,
        .module = shaderModule,
        .pName  = "vertMain",
      };
      vk::PipelineShaderStageCreateInfo fragShaderStageInfo = {
        .stage  = vk::ShaderStageFlagBits::eFragment,
        .module = shaderModule,
        .pName  = "fragMain",
      };

      vk::PipelineShaderStageCreateInfo shaderStages[] = {vertShaderStageInfo, fragShaderStageInfo};

      // Vertex input
      vk::PipelineVertexInputStateCreateInfo vertexInputInfo;

      // Input assembly
      vk::PipelineInputAssemblyStateCreateInfo inputAssembly = {
        .topology = vk::PrimitiveTopology::eTriangleList,
      };

      // Dynamic state
      std::vector<vk::DynamicState> dynamicStates = {
        vk::DynamicState::eViewport,
        vk::DynamicState::eScissor,
      };
      vk::PipelineDynamicStateCreateInfo dynamicState = {
        .dynamicStateCount = static_cast<uint32_t>(dynamicStates.size()),
        .pDynamicStates    = dynamicStates.data(),
      };
      // No need to specify viewport and scissor because they will be specified dynamically
      vk::PipelineViewportStateCreateInfo viewportState = {
        .viewportCount = 1,
        .scissorCount  = 1,
      };

      // Rasterisation
      vk::PipelineRasterizationStateCreateInfo rasterizer = {
        .depthClampEnable        = vk::False,
        .rasterizerDiscardEnable = vk::False,
        .polygonMode             = vk::PolygonMode::eFill,
        .cullMode                = vk::CullModeFlagBits::eBack,
        .frontFace               = vk::FrontFace::eClockwise,
        .depthBiasEnable         = vk::False,
        .depthBiasSlopeFactor    = 1.0f,
        .lineWidth               = 1.0f
      };

      // Multisampling
      vk::PipelineMultisampleStateCreateInfo multisampling = {
        .rasterizationSamples = vk::SampleCountFlagBits::e1,
        .sampleShadingEnable  = vk::False,
      };

      // Color blending
      vk::PipelineColorBlendAttachmentState colorBlendAttachment = {
        .blendEnable    = vk::False,
        .colorWriteMask = vk::ColorComponentFlagBits::eR |
                          vk::ColorComponentFlagBits::eG |
                          vk::ColorComponentFlagBits::eB |
                          vk::ColorComponentFlagBits::eA,
      };
      vk::PipelineColorBlendStateCreateInfo colorBlending = {
        .logicOpEnable   = vk::False,
        .logicOp         = vk::LogicOp::eCopy,
        .attachmentCount = 1,
        .pAttachments    = &colorBlendAttachment,
      };

      // Pipeline layout
      vk::PipelineLayoutCreateInfo pipelineLayoutInfo = {
        .setLayoutCount         = 0,
        .pushConstantRangeCount = 0,
      };
      pipelineLayout = vk::raii::PipelineLayout(device, pipelineLayoutInfo);
    }
    [[nodiscard]] vk::raii::ShaderModule createShaderModule(const std::vector<char> &code) const {
      vk::ShaderModuleCreateInfo createInfo {
        .codeSize = code.size() * sizeof(char),
        .pCode    = reinterpret_cast<const uint32_t *>(code.data()),
      };

      return vk::raii::ShaderModule{device, createInfo};
    }
    vk::SurfaceFormatKHR chooseSwapSurfaceFormat(const std::vector<vk::SurfaceFormatKHR>& availableFormats) {
      for (const auto& availableFormat : availableFormats) {
        if (availableFormat.format == vk::Format::eB8G8R8A8Srgb && availableFormat.colorSpace == vk::ColorSpaceKHR::eSrgbNonlinear) {
          return availableFormat;
        }
      }

      return availableFormats[0];
    }
    vk::PresentModeKHR chooseSwapPresentMode(const std::vector<vk::PresentModeKHR>& availablePresentModes) {
      for (const auto& availablePresentMode : availablePresentModes) {
        if (availablePresentMode == vk::PresentModeKHR::eMailbox) {
          return availablePresentMode;
        }
      }

      return vk::PresentModeKHR::eFifo;
    }
    vk::Extent2D chooseSwapExtent(const vk::SurfaceCapabilitiesKHR& capabilities) {
      if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max()) {
        return capabilities.currentExtent;
      }

      int width, height;
      glfwGetFramebufferSize(window, &width, &height);

      return {
        std::clamp<uint32_t>(width, capabilities.minImageExtent.width, capabilities.maxImageExtent.width),
        std::clamp<uint32_t>(height, capabilities.minImageExtent.height, capabilities.maxImageExtent.height),
      };
    }
    uint32_t findQueueFamilies(vk::raii::PhysicalDevice physicalDevice) {
      // find the index of the first queue family that supports graphics
      std::vector<vk::QueueFamilyProperties> queueFamilyProperties = physicalDevice.getQueueFamilyProperties();

      // get the first index into queueFamilyProperties which both supports graphics and present
      uint32_t queueIndex = ~0;
      for (uint32_t qfpIndex = 0; qfpIndex < queueFamilyProperties.size(); ++qfpIndex) {
        if ((queueFamilyProperties[qfpIndex].queueFlags & vk::QueueFlagBits::eGraphics) &&
            physicalDevice.getSurfaceSupportKHR(qfpIndex, *surface)) {
          queueIndex = qfpIndex;
          break;
        }
      }

      if (queueIndex == ~0) {
        throw std::runtime_error("Could not find a queue for graphics and present -> terminating");
      }

      return queueIndex;
    }

    GLFWwindow                      *window;
    vk::raii::Context                context;
    vk::raii::Instance               instance             = nullptr;
    vk::raii::PhysicalDevice         physicalDevice       = nullptr;
    vk::raii::Device                 device               = nullptr;
    vk::raii::Queue                  graphicsQueue        = nullptr;
    vk::raii::SurfaceKHR             surface              = nullptr;
    vk::raii::SwapchainKHR           swapChain            = nullptr;
    vk::SurfaceFormatKHR             swapChainSurfaceFormat;
    vk::Extent2D                     swapChainExtent;
    vk::Format                       swapChainImageFormat = vk::Format::eUndefined;
    std::vector<vk::Image>           swapChainImages;
    std::vector<vk::raii::ImageView> swapChainImageViews;
    vk::raii::PipelineLayout         pipelineLayout       = nullptr;
};

int main() {
  HelloTriangleApplication app;

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

  return EXIT_SUCCESS;
}