#include <vulkan/vulkan.h>

#include <SDL2/SDL.h>
#include <SDL2/SDL_vulkan.h>

//#define _USE_MATH_DEFINES // Will be needed when/if I need to # include <cmath>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <glm/vec4.hpp>
#include <glm/mat4x4.hpp>

#include <iostream>
#include <vector>
#include <set>
#include <stdexcept>
#include <cstdlib>
#include <optional>
#include <algorithm>

#include "game-gui-sdl.hpp"

using namespace std;
//using namespace glm;

const int SCREEN_WIDTH = 800;
const int SCREEN_HEIGHT = 600;

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

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

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

struct QueueFamilyIndices {
    optional<uint32_t> graphicsFamily;
    optional<uint32_t> presentFamily;

    bool isComplete() {
        return graphicsFamily.has_value() && presentFamily.has_value();
    }
};

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

static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(
      VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
      VkDebugUtilsMessageTypeFlagsEXT messageType,
      const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
      void* pUserData) {
   cerr << "validation layer: " << pCallbackData->pMessage << endl;

   return VK_FALSE;
}

VkResult CreateDebugUtilsMessengerEXT(VkInstance instance,
      const VkDebugUtilsMessengerCreateInfoEXT* pCreateInfo,
      const VkAllocationCallbacks* pAllocator,
      VkDebugUtilsMessengerEXT* pDebugMessenger) {
   auto func = (PFN_vkCreateDebugUtilsMessengerEXT) vkGetInstanceProcAddr(
      instance, "vkCreateDebugUtilsMessengerEXT");

   if (func != nullptr) {
      return func(instance, pCreateInfo, pAllocator, pDebugMessenger);
   } else {
      return VK_ERROR_EXTENSION_NOT_PRESENT;
   }
}

void DestroyDebugUtilsMessengerEXT(VkInstance instance,
      VkDebugUtilsMessengerEXT debugMessenger,
      const VkAllocationCallbacks* pAllocator) {
   auto func = (PFN_vkDestroyDebugUtilsMessengerEXT) vkGetInstanceProcAddr(
      instance, "vkDestroyDebugUtilsMessengerEXT");

   if (func != nullptr) {
      func(instance, debugMessenger, pAllocator);
   }
}

class VulkanGame {
   public:
      void run() {
         if (initWindow() == RTWO_ERROR) {
            return;
         }
         initVulkan();
         mainLoop();
         cleanup();
      }
   private:
      GameGui* gui = new GameGui_SDL();
      SDL_Window* window = nullptr;

      VkInstance instance;
      VkDebugUtilsMessengerEXT debugMessenger;
      VkSurfaceKHR surface;
      SDL_Surface* sdlSurface = nullptr;

      VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;
      VkDevice device;

      VkQueue graphicsQueue;
      VkQueue presentQueue;

      VkSwapchainKHR swapChain;
      vector<VkImage> swapChainImages;
      VkFormat swapChainImageFormat;
      VkExtent2D swapChainExtent;

      vector<VkImageView> swapChainImageViews;

      // both SDL and GLFW create window functions return NULL on failure
      bool initWindow() {
         if (gui->Init() == RTWO_ERROR) {
            cout << "UI library could not be initialized!" << endl;
            return RTWO_ERROR;
         } else {
            // On Apple's OS X you must set the NSHighResolutionCapable Info.plist property to YES,
            // otherwise you will not receive a High DPI OpenGL canvas.

            // TODO: Move this into some generic method in game-gui-sdl
            window = SDL_CreateWindow("Vulkan Game",
               SDL_WINDOWPOS_CENTERED, SDL_WINDOWPOS_CENTERED,
               SCREEN_WIDTH, SCREEN_HEIGHT,
               SDL_WINDOW_VULKAN | SDL_WINDOW_SHOWN);

            if (window == nullptr) {
               cout << "Window could not be created!" << endl;
               return RTWO_ERROR;
            } else {
               return RTWO_SUCCESS;
            }
         }
      }

      void initVulkan() {
         createInstance();
         setupDebugMessenger();
         createSurface();
         pickPhysicalDevice();
         createLogicalDevice();
         createSwapChain();
         createImageViews();
      }

      void createInstance() {
         if (enableValidationLayers && !checkValidationLayerSupport()) {
            throw runtime_error("validation layers requested, but not available!");
         }

         VkApplicationInfo appInfo = {};
         appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
         appInfo.pApplicationName = "Vulkan Game";
         appInfo.applicationVersion = VK_MAKE_VERSION(1, 0, 0);
         appInfo.pEngineName = "No Engine";
         appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);
         appInfo.apiVersion = VK_API_VERSION_1_0;

         VkInstanceCreateInfo createInfo = {};
         createInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
         createInfo.pApplicationInfo = &appInfo;

         vector<const char*> extensions = getRequiredExtensions();
         createInfo.enabledExtensionCount = static_cast<uint32_t>(extensions.size());
         createInfo.ppEnabledExtensionNames = extensions.data();

         cout << endl << "SDL extensions:" << endl;
         for (const char* extensionName : extensions) {
            cout << extensionName << endl;
         }
         cout << endl;

         VkDebugUtilsMessengerCreateInfoEXT debugCreateInfo;
         if (enableValidationLayers) {
            createInfo.enabledLayerCount = static_cast<uint32_t>(validationLayers.size());
            createInfo.ppEnabledLayerNames = validationLayers.data();

            populateDebugMessengerCreateInfo(debugCreateInfo);
            createInfo.pNext = &debugCreateInfo;
         } else {
            createInfo.enabledLayerCount = 0;

            createInfo.pNext = nullptr;
         }

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

      void setupDebugMessenger() {
         if (!enableValidationLayers) return;

         VkDebugUtilsMessengerCreateInfoEXT createInfo;
         populateDebugMessengerCreateInfo(createInfo);

         if (CreateDebugUtilsMessengerEXT(instance, &createInfo, nullptr, &debugMessenger) != VK_SUCCESS) {
            throw runtime_error("failed to setup debug messenger!");
         }
      }

      void createSurface() {
         sdlSurface = SDL_GetWindowSurface(window);

         if (sdlSurface == nullptr) {
             cout << "Could not get SDL Surface! =(" << endl;
         }

         if (!SDL_Vulkan_CreateSurface(window, instance, &surface)) {
            throw runtime_error("failed to create window surface!");
         }

         /*
         if (glfwCreateWindowSurface(instance, window, nullptr, &surface) != VK_SUCCESS) {
            throw runtime_error("failed to create window surface!");
         }
         */
      }

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

         if (deviceCount == 0) {
            throw runtime_error("failed to find GPUs with Vulkan support!");
         }

         vector<VkPhysicalDevice> devices(deviceCount);
         vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());

         cout << endl << "Graphics cards:" << endl;
         for (const VkPhysicalDevice& device : devices) {
            if (isDeviceSuitable(device)) {
               physicalDevice = device;
               break;
            }
         }
         cout << endl;

         if (physicalDevice == VK_NULL_HANDLE) {
            throw runtime_error("failed to find a suitable GPU!");
         }
      }

      bool isDeviceSuitable(VkPhysicalDevice device) {
         VkPhysicalDeviceProperties deviceProperties;
         VkPhysicalDeviceFeatures deviceFeatures;

         vkGetPhysicalDeviceProperties(device, &deviceProperties);
         vkGetPhysicalDeviceFeatures(device, &deviceFeatures);

         cout << "Device: " << deviceProperties.deviceName << endl;

         QueueFamilyIndices indices = findQueueFamilies(device);

         bool extensionsSupported = checkDeviceExtensionSupport(device);

         bool swapChainAdequate = false;

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

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

      bool checkDeviceExtensionSupport(VkPhysicalDevice device) {
         uint32_t extensionCount;
         vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, nullptr);

         vector<VkExtensionProperties> availableExtensions(extensionCount);
         vkEnumerateDeviceExtensionProperties(device, nullptr, &extensionCount, availableExtensions.data());

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

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

         return requiredExtensions.empty();
      }

      void createLogicalDevice() {
         QueueFamilyIndices indices = findQueueFamilies(physicalDevice);

         vector<VkDeviceQueueCreateInfo> queueCreateInfos;
         set<uint32_t> uniqueQueueFamilies = {indices.graphicsFamily.value(), indices.presentFamily.value()};

         float queuePriority = 1.0f;
         for (uint32_t queueFamily : uniqueQueueFamilies) {
            VkDeviceQueueCreateInfo queueCreateInfo = {};

            queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
            queueCreateInfo.queueFamilyIndex = queueFamily;
            queueCreateInfo.queueCount = 1;
            queueCreateInfo.pQueuePriorities = &queuePriority;

            queueCreateInfos.push_back(queueCreateInfo);
         }

         VkPhysicalDeviceFeatures deviceFeatures = {};

         VkDeviceCreateInfo createInfo = {};
         createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;

         createInfo.queueCreateInfoCount = static_cast<uint32_t>(queueCreateInfos.size());;
         createInfo.pQueueCreateInfos = queueCreateInfos.data();

         createInfo.pEnabledFeatures = &deviceFeatures;

         createInfo.enabledExtensionCount = static_cast<uint32_t>(deviceExtensions.size());
         createInfo.ppEnabledExtensionNames = deviceExtensions.data();

         // These fields are ignored  by up-to-date Vulkan implementations,
         // but it's a good idea to set them for backwards compatibility
         if (enableValidationLayers) {
            createInfo.enabledLayerCount = static_cast<uint32_t>(validationLayers.size());
            createInfo.ppEnabledLayerNames = validationLayers.data();
         } else {
            createInfo.enabledLayerCount = 0;
         }

         if (vkCreateDevice(physicalDevice, &createInfo, nullptr, &device) != VK_SUCCESS) {
            throw runtime_error("failed to create logical device!");
         }

         vkGetDeviceQueue(device, indices.graphicsFamily.value(), 0, &graphicsQueue);
         vkGetDeviceQueue(device, indices.presentFamily.value(), 0, &presentQueue);
      }

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

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

         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;
      }

      QueueFamilyIndices findQueueFamilies(VkPhysicalDevice device) {
         QueueFamilyIndices indices;

         uint32_t queueFamilyCount = 0;
         vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);

         vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
         vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

         int i = 0;
         for (const auto& queueFamily : queueFamilies) {
            if (queueFamily.queueCount > 0 && queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT) {
               indices.graphicsFamily = i;
            }

            VkBool32 presentSupport = false;
            vkGetPhysicalDeviceSurfaceSupportKHR(device, i, surface, &presentSupport);

            if (queueFamily.queueCount > 0 && presentSupport) {
               indices.presentFamily = i;
            }

            if (indices.isComplete()) {
               break;
            }

            i++;
         }

         return indices;
      }

      SwapChainSupportDetails querySwapChainSupport(VkPhysicalDevice device) {
         SwapChainSupportDetails details;

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

         uint32_t formatCount;
         vkGetPhysicalDeviceSurfaceFormatsKHR(device, surface, &formatCount, nullptr);

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

         uint32_t presentModeCount;
         vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, nullptr);

         if (presentModeCount != 0) {
            details.presentModes.resize(presentModeCount);
            vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, details.presentModes.data());
         }

         return details;
      }

      VkSurfaceFormatKHR chooseSwapSurfaceFormat(const vector<VkSurfaceFormatKHR>& availableFormats) {
         for (const auto& availableFormat : availableFormats) {
            if (availableFormat.format == VK_FORMAT_B8G8R8A8_UNORM && availableFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
               return availableFormat;
            }
         }

         return availableFormats[0];
      }

      VkPresentModeKHR chooseSwapPresentMode(const vector<VkPresentModeKHR>& availablePresentModes) {
         VkPresentModeKHR bestMode = VK_PRESENT_MODE_FIFO_KHR;

         for (const auto& availablePresentMode : availablePresentModes) {
            if (availablePresentMode == VK_PRESENT_MODE_MAILBOX_KHR) {
               return availablePresentMode;
            } else if (availablePresentMode == VK_PRESENT_MODE_IMMEDIATE_KHR) {
               bestMode = availablePresentMode;
            }
         }

         return bestMode;
      }

      VkExtent2D chooseSwapExtent(const VkSurfaceCapabilitiesKHR& capabilities) {
         if (capabilities.currentExtent.width != numeric_limits<uint32_t>::max()) {
            return capabilities.currentExtent;
         } else {
            VkExtent2D actualExtent = { SCREEN_WIDTH, SCREEN_HEIGHT };

            actualExtent.width = max(capabilities.minImageExtent.width, min(capabilities.maxImageExtent.width, actualExtent.width));
            actualExtent.height = max(capabilities.minImageExtent.height, min(capabilities.maxImageExtent.height, actualExtent.height));

            return actualExtent;
         }
      }

      void populateDebugMessengerCreateInfo(VkDebugUtilsMessengerCreateInfoEXT& createInfo) {
         createInfo = {};
         createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
         createInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
         createInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
         createInfo.pfnUserCallback = debugCallback;
      }

      vector<const char*> getRequiredExtensions() {
         uint32_t extensionCount = 0;
         SDL_Vulkan_GetInstanceExtensions(window, &extensionCount, nullptr);

         vector<const char*> extensions(extensionCount);
         SDL_Vulkan_GetInstanceExtensions(window, &extensionCount, extensions.data());

         if (enableValidationLayers) {
            extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
         }

         return extensions;
      }

      void createSwapChain() {
         SwapChainSupportDetails swapChainSupport = querySwapChainSupport(physicalDevice);

         VkSurfaceFormatKHR surfaceFormat = chooseSwapSurfaceFormat(swapChainSupport.formats);
         VkPresentModeKHR presentMode = chooseSwapPresentMode(swapChainSupport.presentModes);
         VkExtent2D extent = chooseSwapExtent(swapChainSupport.capabilities);

         uint32_t imageCount = swapChainSupport.capabilities.minImageCount + 1;
         if (swapChainSupport.capabilities.maxImageCount > 0 &&
               imageCount > swapChainSupport.capabilities.maxImageCount) {
            imageCount = swapChainSupport.capabilities.maxImageCount;
         }

         VkSwapchainCreateInfoKHR createInfo = {};

         createInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
         createInfo.surface = surface;
         createInfo.minImageCount = imageCount;
         createInfo.imageFormat = surfaceFormat.format;
         createInfo.imageColorSpace = surfaceFormat.colorSpace;
         createInfo.imageExtent = extent;
         createInfo.imageArrayLayers = 1;
         createInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

         QueueFamilyIndices indices = findQueueFamilies(physicalDevice);
         uint32_t queueFamilyIndices[] = {indices.graphicsFamily.value(), indices.presentFamily.value()};

         if (indices.graphicsFamily != indices.presentFamily) {
            createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
            createInfo.queueFamilyIndexCount = 2;
            createInfo.pQueueFamilyIndices = queueFamilyIndices;
         } else {
            createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
            createInfo.queueFamilyIndexCount = 0; // Optional
            createInfo.pQueueFamilyIndices = nullptr;
         }

         createInfo.preTransform = swapChainSupport.capabilities.currentTransform;
         createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
         createInfo.presentMode = presentMode;
         createInfo.clipped = VK_TRUE;
         createInfo.oldSwapchain = VK_NULL_HANDLE;

         if (vkCreateSwapchainKHR(device, &createInfo, nullptr, &swapChain) != VK_SUCCESS) {
            throw runtime_error("failed to create swap chain!");
         }

         vkGetSwapchainImagesKHR(device, swapChain, &imageCount, nullptr);
         swapChainImages.resize(imageCount);
         vkGetSwapchainImagesKHR(device, swapChain, &imageCount, swapChainImages.data());

         swapChainImageFormat = surfaceFormat.format;
         swapChainExtent = extent;
      }

      void createImageViews() {
         swapChainImageViews.resize(swapChainImages.size());

         for (size_t i=0; i<swapChainImages.size(); i++) {
            VkImageViewCreateInfo createInfo = {};
            createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
            createInfo.image = swapChainImages[i];

            createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
            createInfo.format = swapChainImageFormat;

            createInfo.components.r = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.g = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.b = VK_COMPONENT_SWIZZLE_IDENTITY;
            createInfo.components.a = VK_COMPONENT_SWIZZLE_IDENTITY;

            createInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
            createInfo.subresourceRange.baseMipLevel = 0;
            createInfo.subresourceRange.levelCount = 1;
            createInfo.subresourceRange.baseArrayLayer = 0;
            createInfo.subresourceRange.layerCount = 1;

            if (vkCreateImageView(device, &createInfo, nullptr, &swapChainImageViews[i]) != VK_SUCCESS) {
               throw runtime_error("failed to create image views!");
            }
         }
      }

      void mainLoop() {
         // TODO: Create some generic event-handling functions in game-gui-*
         SDL_Event e;
         bool quit = false;

         while (!quit) {
            while (SDL_PollEvent(&e)) {
               if (e.type == SDL_QUIT) {
                  quit = true;
               }
               if (e.type == SDL_KEYDOWN) {
                  quit = true;
               }
               if (e.type == SDL_MOUSEBUTTONDOWN) {
                  quit = true;
               }

               /*
               SDL_FillRect(sdlSurface, nullptr, SDL_MapRGB(sdlSurface->format, 0xFF, 0xFF, 0xFF));

               SDL_UpdateWindowSurface(window);
               */
            }
         }
      }

      void cleanup() {
         for (auto imageView : swapChainImageViews) {
            vkDestroyImageView(device, imageView, nullptr);
         }

         vkDestroySwapchainKHR(device, swapChain, nullptr);
         vkDestroyDevice(device, nullptr);

         if (enableValidationLayers) {
            DestroyDebugUtilsMessengerEXT(instance, debugMessenger, nullptr);
         }

         vkDestroySurfaceKHR(instance, surface, nullptr);
         vkDestroyInstance(instance, nullptr);

         // TODO: Move this into some generic method in game-gui-sdl
         SDL_DestroyWindow(window);

         gui->Shutdown();
         delete gui;
      }
};

int main(int argc, char* argv[]) {

#ifdef NDEBUG
   cout << "DEBUGGING IS OFF" << endl;
#else
   cout << "DEBUGGING IS ON" << endl;
#endif

   glm::mat4 matrix;
   glm::vec4 vec;
   glm::vec4 test = matrix * vec;

   cout << "Starting Vulkan game..." << endl;

   VulkanGame game;

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

   cout << "Finished running the game" << endl;

   return EXIT_SUCCESS;
}