#define GLFW_INCLUDE_VULKAN

#include <cstring>
#include "GLFW/glfw3.h"

#include "iostream"
#include "stdexcept"
#include "cstdlib"
#include "vector"

const std::vector<const char*> deviceExtensions = {
  VK_KHR_SWAPCHAIN_EXTENSION_NAME,
};

const std::vector<const char*> validationLayers = {
#ifdef NDEBUG
#else
    "VK_LAYER_KHRONOS_validation"
#endif
};

struct SwapChainSupportDetails {
  VkSurfaceCapabilitiesKHR capabilities {};
  std::vector<VkSurfaceFormatKHR> formats;
  std::vector<VkPresentModeKHR> presentModes;
};

class Application {
public:
  void run() {
    initWindow();
    initVulkan();
    mainLoop();
    cleanup();
  }

private:
  void initWindow() {
    glfwInit();

    glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
    glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
    mWindow = glfwCreateWindow(800, 600, "App", nullptr, nullptr);
  }

  void initVulkan() {
    createInstance();
    createWindowSurface();
    pickPhysicalDevice();
    findPhysicalDeviceQueueFamilies();
    createLogicalDevice();
    getQueues();
    createSwapChain();
  }

  void mainLoop() {
    while(!glfwWindowShouldClose(mWindow)) {
      glfwPollEvents();
    }
  }

  void cleanup() {
    vkDestroySwapchainKHR(mDevice, mSwapChain, nullptr);
    vkDestroyDevice(mDevice, nullptr);
    vkDestroySurfaceKHR(mInstance, mSurface, nullptr);
    vkDestroyInstance(mInstance, nullptr);
    glfwDestroyWindow(mWindow);
    glfwTerminate();
  }


  void createInstance() {
    if (!checkValidationLayerSupport()) {
      throw std::runtime_error("no required validation layers present");
    }

    uint32_t glfwExtensionCount = 0;
    const char** glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);

    VkApplicationInfo appInfo {
      .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
      .pApplicationName = "App",
      .applicationVersion = VK_MAKE_VERSION(0, 0, 0),
      .pEngineName = "no",
      .engineVersion = VK_MAKE_VERSION(0, 0, 0),
      .apiVersion = VK_API_VERSION_1_0,
    };

    VkInstanceCreateInfo createInfo {
      .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
      .pApplicationInfo = &appInfo,
      .enabledLayerCount = (uint32_t) validationLayers.size(),
      .ppEnabledLayerNames = validationLayers.data(),
      .enabledExtensionCount = glfwExtensionCount,
      .ppEnabledExtensionNames = glfwExtensions,
    };

    if (vkCreateInstance(&createInfo, nullptr, &mInstance) != VK_SUCCESS) {
      throw std::runtime_error("failed to create instance!");
    }
  }

  static bool checkValidationLayerSupport() {
    uint32_t layerCount;
    vkEnumerateInstanceLayerProperties(&layerCount, nullptr);

    std::vector<VkLayerProperties> availableLayers(layerCount);
    vkEnumerateInstanceLayerProperties(&layerCount, availableLayers.data());

    for (const auto& requestedLayer : validationLayers) {
      bool presents = false;
      for (auto &layer: availableLayers) {
        if (strcmp(requestedLayer, layer.layerName) == 0) {
          presents = true;
          break;
        }
      }

      if (!presents) return false;
    }

    return true;
  }

  void pickPhysicalDevice() {
    uint32_t deviceCount = 0;
    vkEnumeratePhysicalDevices(mInstance, &deviceCount, nullptr);

    if (deviceCount == 0) {
      throw std::runtime_error("no gpu with vulkan support");
    }

    std::vector<VkPhysicalDevice> devices(deviceCount);
    vkEnumeratePhysicalDevices(mInstance, &deviceCount, devices.data());

    for (const auto& device : devices) {
      if (isDeviceSuitable(device)) {
        mPhysicalDevice = device;
        break;
      }
    }

    if (!mPhysicalDevice) throw std::runtime_error("no suitable gpu");
  }

  bool isDeviceSuitable(VkPhysicalDevice device) {
    VkPhysicalDeviceProperties properties;
    VkPhysicalDeviceFeatures features;

    vkGetPhysicalDeviceProperties(device, &properties);
    vkGetPhysicalDeviceFeatures(device, &features);

    // if (properties.deviceType != VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) return false;
    if (!features.geometryShader) return false;
    if (!checkDeviceExtensions(device)) return false;
    if (!checkDeviceSwapChain(device)) return false;

    return true;
  }

  bool checkDeviceSwapChain(VkPhysicalDevice device) {
    SwapChainSupportDetails details = querySwapChainSupportDetails(device);
    return !(details.presentModes.empty() || details.formats.empty());
  }

  static bool checkDeviceExtensions(VkPhysicalDevice device) {
    uint32_t extensionsCount = 0;
    vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionsCount, nullptr);

    std::vector<VkExtensionProperties> extensions(extensionsCount);
    vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionsCount, extensions.data());

    for (const auto& requiredExtension : deviceExtensions) {
      bool found = false;
      for (const auto& extension : extensions) {
        if (strcmp(requiredExtension, extension.extensionName) == 0) {
          found = true;
          break;
        }
      }
      if (!found) return false;
    }

    return true;
  }

  void findPhysicalDeviceQueueFamilies() {
    uint32_t queuesCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(mPhysicalDevice, &queuesCount, nullptr);

    std::vector<VkQueueFamilyProperties> queueFamilyProperties(queuesCount);
    vkGetPhysicalDeviceQueueFamilyProperties(mPhysicalDevice, &queuesCount, queueFamilyProperties.data());

    int index = 0;
    for (const auto& familyProperty : queueFamilyProperties) {
      if (familyProperty.queueFlags & VK_QUEUE_GRAPHICS_BIT) {
        mGraphicsQueueFamilyIndex = index;
      }

      VkBool32 presentationQueueSupport = false;
      vkGetPhysicalDeviceSurfaceSupportKHR(mPhysicalDevice, index, mSurface, &presentationQueueSupport);

      if (presentationQueueSupport) {
        mPresentationQueueFamilyIndex = index;
      }

      index++;
    }

    if (mPresentationQueueFamilyIndex == -1 || mGraphicsQueueFamilyIndex == -1) {
      throw std::runtime_error("nu require queue families found");
    }
  }

  SwapChainSupportDetails querySwapChainSupportDetails(VkPhysicalDevice device) {
    SwapChainSupportDetails details;

    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, mSurface, &details.capabilities);

    uint32_t formatCount = 0;
    vkGetPhysicalDeviceSurfaceFormatsKHR(device, mSurface, &formatCount, nullptr);
    if (formatCount != 0) {
      details.formats.resize(formatCount);
      vkGetPhysicalDeviceSurfaceFormatsKHR(device, mSurface, &formatCount, details.formats.data());
    }

    uint32_t modeCount = 0;
    vkGetPhysicalDeviceSurfacePresentModesKHR(device, mSurface, &modeCount, nullptr);
    if (modeCount != 0) {
      details.presentModes.resize(modeCount);
      vkGetPhysicalDeviceSurfacePresentModesKHR(device, mSurface, &modeCount, details.presentModes.data());
    }

    return details;
  }

  void createLogicalDevice() {
    float queuePriority = 1.f;

    VkDeviceQueueCreateInfo graphicsQueueCreateInfos {
      .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
      .queueFamilyIndex = mGraphicsQueueFamilyIndex,
      .queueCount = 1,
      .pQueuePriorities = &queuePriority,
    };

    VkDeviceQueueCreateInfo presentationQueueCreateInfos {
        .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
        .queueFamilyIndex = mPresentationQueueFamilyIndex,
        .queueCount = 1,
        .pQueuePriorities = &queuePriority,
    };

    std::vector<VkDeviceQueueCreateInfo> queues = { graphicsQueueCreateInfos };

    if (mPresentationQueueFamilyIndex != mGraphicsQueueFamilyIndex) {
      queues.push_back(presentationQueueCreateInfos);
    }

    VkPhysicalDeviceFeatures features {};

    VkDeviceCreateInfo deviceCreateInfo {
      .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
      .queueCreateInfoCount = (uint32_t) queues.size(),
      .pQueueCreateInfos = queues.data(),
      .enabledExtensionCount = (uint32_t) deviceExtensions.size(),
      .ppEnabledExtensionNames = deviceExtensions.data(),
      .pEnabledFeatures = &features,
    };

    vkCreateDevice(mPhysicalDevice, &deviceCreateInfo, nullptr, &mDevice);

    if (mDevice == VK_NULL_HANDLE) throw std::runtime_error("failed to create vulkan logical device");
  }

  void getQueues() {
    vkGetDeviceQueue(mDevice, mGraphicsQueueFamilyIndex, 0, &mGraphicsQueue);
    vkGetDeviceQueue(mDevice, mPresentationQueueFamilyIndex, 0, &mPresentQueue);
  }


  void createWindowSurface() {
    if (glfwCreateWindowSurface(mInstance, mWindow, nullptr, &mSurface) != VK_SUCCESS) {
      throw std::runtime_error("failed to create vulkan window surface");
    }
  }

  void createSwapChain() {
    SwapChainSupportDetails details = querySwapChainSupportDetails(mPhysicalDevice);

    VkSurfaceFormatKHR surfaceFormat = pickSwapChainSurfaceFormat(details);
    VkPresentModeKHR presentMode = pickSwapChainPresentMode(details);
    VkExtent2D extent2D = pickSwapChainExtent(details.capabilities);

    // +1 so will not have to wait for device to finish frame to query an image to render to
    uint32_t imageCount = details.capabilities.minImageCount + 1;

    if (details.capabilities.maxImageCount > 0 && imageCount > details.capabilities.maxImageCount) {
      imageCount = details.capabilities.maxImageCount;
    }

    VkSwapchainCreateInfoKHR createInfo {
      .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
      .surface = mSurface,
      .minImageCount = imageCount,
      .imageFormat = surfaceFormat.format,
      .imageExtent = extent2D,
      .imageArrayLayers = 1,
      .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
      .preTransform = details.capabilities.currentTransform,
      .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
      .presentMode = presentMode,
      .clipped = VK_TRUE,
      .oldSwapchain = VK_NULL_HANDLE,
    };

    uint32_t queueIndices[] = { mGraphicsQueueFamilyIndex, mPresentationQueueFamilyIndex };

    if (mGraphicsQueueFamilyIndex != mPresentationQueueFamilyIndex) {
      createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
      createInfo.queueFamilyIndexCount = 2;
      createInfo.pQueueFamilyIndices = queueIndices;
    } else {
      createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    }

    if (vkCreateSwapchainKHR(mDevice, &createInfo, nullptr, &mSwapChain) != VK_SUCCESS) {
      throw std::runtime_error("failed to create swap chain");
    }

    mSwapChainExtent = extent2D;
    mSwapChainFormat = surfaceFormat.format;

    uint32_t createdImageCount = 0;
    vkGetSwapchainImagesKHR(mDevice, mSwapChain, &createdImageCount, nullptr);
    mSwapChainImages.resize(createdImageCount);
    vkGetSwapchainImagesKHR(mDevice, mSwapChain, &createdImageCount, mSwapChainImages.data());
  }

  static VkSurfaceFormatKHR pickSwapChainSurfaceFormat(const SwapChainSupportDetails& details) {
    for (const auto& format : details.formats) {
      if (format.format == VK_FORMAT_B8G8R8A8_SRGB && format.colorSpace == VK_COLORSPACE_SRGB_NONLINEAR_KHR) {
        return format;
      }
    }

    return details.formats.front();
  }

  static VkPresentModeKHR pickSwapChainPresentMode(const SwapChainSupportDetails& details) {
    for (const auto& mode : details.presentModes) {
      if (mode == VK_PRESENT_MODE_MAILBOX_KHR) {
        return mode;
      }
    }

    return VK_PRESENT_MODE_FIFO_KHR; // guaranteed exists
  }

  VkExtent2D pickSwapChainExtent(const VkSurfaceCapabilitiesKHR& capabilities) {
    // if set by vulkan just keep it
    if (capabilities.currentExtent.width != std::numeric_limits<uint32_t>::max()) {
      return capabilities.currentExtent;
    }

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

    VkExtent2D out = { (uint32_t) width, (uint32_t) height };

    out.width = std::clamp(out.width, capabilities.minImageExtent.width, capabilities.maxImageExtent.width);
    out.height = std::clamp(out.height, capabilities.minImageExtent.height, capabilities.maxImageExtent.height);

    return out;
  }

private:
  GLFWwindow* mWindow = nullptr;

  uint32_t mGraphicsQueueFamilyIndex = -1;
  uint32_t mPresentationQueueFamilyIndex = -1;

  VkInstance mInstance = VK_NULL_HANDLE;
  VkPhysicalDevice mPhysicalDevice = VK_NULL_HANDLE;
  VkDevice mDevice = VK_NULL_HANDLE;

  VkQueue mGraphicsQueue = VK_NULL_HANDLE;
  VkQueue mPresentQueue = VK_NULL_HANDLE;

  VkSurfaceKHR mSurface = VK_NULL_HANDLE;

  VkSwapchainKHR mSwapChain = VK_NULL_HANDLE;
  std::vector<VkImage> mSwapChainImages;
  VkFormat mSwapChainFormat {};
  VkExtent2D mSwapChainExtent {};
};

int main() {
  Application app;

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

  return EXIT_SUCCESS;
}