#include <fmt/format.h>
#include <vector>
#include <stdexcept>
#include <string.h>
#include <iostream>
#include <map>
#include <set>

#include "aethiopicus.hpp"
#include "vkutil.hpp"

#include "enginever.hpp"

std::vector<int> split(std::string phrase, std::string delimiter){
	std::vector<int> list;
	size_t pos = 0;
	std::string token;
	while ((pos = phrase.find(delimiter)) != std::string::npos) {
		token = phrase.substr(0, pos);
		list.push_back(stoi(token));
		phrase.erase(0, pos + delimiter.length());
	}

#ifndef NDEBUG
	if ((pos = phrase.find("-")) != std::string::npos) {
		phrase = phrase.substr(0, 1);
	}
#endif

	list.push_back(stoi(phrase));
	return list;
}

VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(
	VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
	VkDebugUtilsMessageTypeFlagsEXT messageType,
	const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
	void* pUserData) {

	if (hibis::aethiopicus::Aethiopicus* renderer = static_cast<hibis::aethiopicus::Aethiopicus*>(pUserData)) {
		LoggingSeverity severity;

		switch (messageSeverity) {
			case VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT:
				severity = Message;
				break;
			case VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT:
				severity = Information;
				break;
			case VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT:
				severity = Warning;
				break;
			case VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT:
				severity = Error;
				break;
			case VK_DEBUG_UTILS_MESSAGE_SEVERITY_FLAG_BITS_MAX_ENUM_EXT:
				severity = Message;
				break;
		}

		renderer->mLogger(severity, "VULKAN: " + (std::string)(pCallbackData->pMessage));
	} else std::cout << "pUserData is not equivalent to hibis::aethiopicus::Aethiopicus." << std::endl;

	return VK_FALSE;
}

void keyCallback(GLFWwindow* window, int key, int scancode, int action, int mods) {
	if (hibis::aethiopicus::Aethiopicus* renderer = static_cast<hibis::aethiopicus::Aethiopicus*>(glfwGetWindowUserPointer(window))) {
		// Prepare type
		hibis::PressTypes type;
		switch (action) {
			case GLFW_PRESS:
				type = hibis::Pressed;
				break;
			case GLFW_REPEAT:
				type = hibis::Held;
				break;
			case GLFW_RELEASE:
				type = hibis::Released;
				break;
			default:
				type = hibis::Unknown;
				break;
		}

		// Tell input manager the new state
		renderer->mInputManager.setKeyPressState(scancode, type);
	}
}

namespace hibis::aethiopicus {
	Aethiopicus::Aethiopicus(std::string title, std::string internalName, std::string version, IntVec2 size, LoggerCallback logger, bool useValidationLayers) : Renderer(logger), mUseValidationLayers(useValidationLayers) {
		// GLFW Window
		glfwInit();

		glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
		glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
		mWindow = glfwCreateWindow(size.x, size.y, title.c_str(), nullptr, nullptr);

		if (mWindow == NULL) throw std::runtime_error("Couldn't create a window");

		glfwSetWindowUserPointer(mWindow, this);
		glfwSetKeyCallback(mWindow, keyCallback);

		std::vector<int> hibisVerNums = split(HIBIS_VERSION, ".");
		std::vector<int> appVerNums = split(version, ".");

		VkApplicationInfo appInfo {};
		appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
		appInfo.pApplicationName = internalName.c_str();
		appInfo.applicationVersion = VK_MAKE_VERSION(appVerNums[0], appVerNums[1], appVerNums[2]);
		appInfo.pEngineName = "Hibis";
		appInfo.engineVersion = VK_MAKE_VERSION(hibisVerNums[0], hibisVerNums[1], hibisVerNums[2]);
		appInfo.apiVersion = VK_API_VERSION_1_3;

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

		mGLFWExtensions = glfwGetRequiredInstanceExtensions(&mGLFWExtensionCount);

		for(uint32_t i = 0; i < mGLFWExtensionCount; i++) {
			mGLFWExtensionVector.emplace_back(mGLFWExtensions[i]);
		}

		mGLFWExtensionVector.emplace_back(VK_KHR_PORTABILITY_ENUMERATION_EXTENSION_NAME);

		if (useValidationLayers) {
			mGLFWExtensionVector.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
		}

		if (useValidationLayers && !checkValidationLayerSupport()) {
			throw std::runtime_error("validation layers requested, but not available!");
		}

		instanceCreateInfo.flags |= VK_INSTANCE_CREATE_ENUMERATE_PORTABILITY_BIT_KHR;

		instanceCreateInfo.enabledExtensionCount = (uint32_t) mGLFWExtensionVector.size();
		instanceCreateInfo.ppEnabledExtensionNames = mGLFWExtensionVector.data();

		uint32_t extensionCount = 0;
		vkEnumerateInstanceExtensionProperties(nullptr, &extensionCount, nullptr);

		mVulkanExtensions.resize(extensionCount);

		vkEnumerateInstanceExtensionProperties(nullptr, &extensionCount, mVulkanExtensions.data());

		if (useValidationLayers) {
			instanceCreateInfo.enabledLayerCount = static_cast<uint32_t>(mValidationLayers.size());
			instanceCreateInfo.ppEnabledLayerNames = mValidationLayers.data();
		} else {
			instanceCreateInfo.enabledLayerCount = 0;
		}

		VkResult result = vkCreateInstance(&instanceCreateInfo, nullptr, &mVulkanInstance);

		if (useValidationLayers && !checkValidationLayerSupport()) {
			throw std::runtime_error("validation layers requested, but not available!");
		}

		if (result != VK_SUCCESS) {
			logger(Fatal, "Could not create an instance of Vulkan");
			throw std::runtime_error("aethiopicus launch err, check logs");
		}

		if (glfwCreateWindowSurface(mVulkanInstance, mWindow, nullptr, &mVulkanSurface) != VK_SUCCESS) throw std::runtime_error("Couldn't make a surface for vulkan to use");

		if (useValidationLayers) {
			VkDebugUtilsMessengerCreateInfoEXT 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;
			createInfo.pUserData = (void*)this;

			// Remove vulkan warnings
			createInfo.pNext = NULL;
			createInfo.flags = 0;

			if (createDebugUtilsMessengerEXT(mVulkanInstance, &createInfo, nullptr, &mDebugMessenger) != VK_SUCCESS) {
				throw std::runtime_error("failed to set up debug messenger!");
			}
		}

		uint32_t deviceCount = 0;
		vkEnumeratePhysicalDevices(mVulkanInstance, &deviceCount, nullptr);

		if (deviceCount == 0) throw std::runtime_error("No vulkan support on available graphics cards");

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

		std::multimap<int, VkPhysicalDevice> candidates;

		for (const auto& device : devices) {
			// Get properties here for logging purposes
			VkPhysicalDeviceProperties deviceProperties;
			vkGetPhysicalDeviceProperties(device, &deviceProperties);

			int score = rateDeviceSuitability(device, deviceProperties);

			std::string deviceTypeAsString;

			switch (deviceProperties.deviceType) {
				case VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU";
					break;
				case VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU";
					break;
				case VK_PHYSICAL_DEVICE_TYPE_CPU:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_CPU";
					break;
				case VK_PHYSICAL_DEVICE_TYPE_MAX_ENUM:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_MAX_ENUM";
					break;
				case VK_PHYSICAL_DEVICE_TYPE_OTHER:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_OTHER";
					break;
				case VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU:
					deviceTypeAsString = "VK_PHYSICAL_DEVICE_TYPE_VIRTUAL_GPU";
					break;
			}

			// Logging the info
			mLogger(Information, fmt::format("Device {} has a score of {}", deviceProperties.deviceName, score));
			mLogger(Information, fmt::format("{}'s device type: {}", deviceProperties.deviceName, deviceTypeAsString));

			// Insert into candidates map
			candidates.insert(std::make_pair(score, device));
		}

		// Validate that the first candidate has a better score than 0.
		if (candidates.rbegin()->first > 0) {
			mVulkanGraphicsDevice = candidates.rbegin()->second;
		} else {
			throw std::runtime_error("failed to find a suitable GPU!");
		}

		// Logical device creation
		QueueFamilyIndicies indices = findQueueFamilies(mVulkanGraphicsDevice, mVulkanSurface);

		std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;
		std::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 logicalDeviceCreateInfo{};
		logicalDeviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;

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

		logicalDeviceCreateInfo.pEnabledFeatures = &deviceFeatures;

		logicalDeviceCreateInfo.enabledExtensionCount = static_cast<uint32_t>(mRequiredDeviceExtensions.size());
		logicalDeviceCreateInfo.ppEnabledExtensionNames = mRequiredDeviceExtensions.data();

		if (mUseValidationLayers) {
			logicalDeviceCreateInfo.enabledLayerCount = static_cast<uint32_t>(mValidationLayers.size());
			logicalDeviceCreateInfo.ppEnabledLayerNames = mValidationLayers.data();
		} else {
			logicalDeviceCreateInfo.enabledLayerCount = 0;
		}

		if (vkCreateDevice(mVulkanGraphicsDevice, &logicalDeviceCreateInfo, nullptr, &mVulkanLogicalDevice) != VK_SUCCESS) {
			throw std::runtime_error("failed to create logical device!");
		}

		vkGetDeviceQueue(mVulkanLogicalDevice, indices.graphicsFamily.value(), 0, &mVulkanGraphicsQueue);
		vkGetDeviceQueue(mVulkanLogicalDevice, indices.presentFamily.value(), 0, &mVulkanPresentQueue);

		// Create a swap chain
		SwapChainSupportDetails swapChainSupport = querySwapChainSupport(mVulkanGraphicsDevice, mVulkanSurface);

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

		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 = mVulkanSurface;

		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;

		QueueFamilyIndicies graphicsIndicies = findQueueFamilies(mVulkanGraphicsDevice, mVulkanSurface);
		uint32_t queueFamilyIndices[] = {graphicsIndicies.graphicsFamily.value(), graphicsIndicies.presentFamily.value()};

		if (graphicsIndicies.graphicsFamily != graphicsIndicies.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; // Optional
		}

		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(mVulkanLogicalDevice, &createInfo, nullptr, &mVulkanSwapchain) != VK_SUCCESS) {
			throw std::runtime_error("failed to create swap chain!");
		}

		mSwapChainImageFormat = surfaceFormat.format;
		mSwapChainExtent = extent;

		vkGetSwapchainImagesKHR(mVulkanLogicalDevice, mVulkanSwapchain, &imageCount, nullptr);
		mSwapchainImages.resize(imageCount);
		vkGetSwapchainImagesKHR(mVulkanLogicalDevice, mVulkanSwapchain, &imageCount, mSwapchainImages.data());

		// We're still fucking going with initialization, fucking christ
		mSwapChainImageViews.resize(mSwapchainImages.size());
		for (size_t i = 0; i < mSwapchainImages.size(); i++) {
			VkImageViewCreateInfo createInfo{};
			createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
			createInfo.image = mSwapchainImages[i];
			createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
			createInfo.format = mSwapChainImageFormat;

			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(mVulkanLogicalDevice, &createInfo, nullptr, &mSwapChainImageViews[i]) != VK_SUCCESS) {
				throw std::runtime_error("failed to create image views!");
			}
		}
	}

	Aethiopicus::~Aethiopicus() {
		if (mUseValidationLayers) {
			destroyDebugUtilsMessengerEXT(mVulkanInstance, mDebugMessenger, nullptr);
		}

		for (auto imageView : mSwapChainImageViews) {
			vkDestroyImageView(mVulkanLogicalDevice, imageView, nullptr);
		}

		vkDestroySwapchainKHR(mVulkanLogicalDevice, mVulkanSwapchain, nullptr);
		vkDestroyDevice(mVulkanLogicalDevice, nullptr);
		vkDestroySurfaceKHR(mVulkanInstance, mVulkanSurface, nullptr);
		vkDestroyInstance(mVulkanInstance, nullptr);
		glfwDestroyWindow(mWindow);
		glfwTerminate();
	}

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

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

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

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

			if (!layerFound) {
				return false;
			}
		}

		return true;
	}

	int Aethiopicus::rateDeviceSuitability(VkPhysicalDevice device, VkPhysicalDeviceProperties deviceProperties) {
		// TODO: Better scoring, accounting for 3D limits
		VkPhysicalDeviceFeatures deviceFeatures;
		vkGetPhysicalDeviceFeatures(device, &deviceFeatures);

		QueueFamilyIndicies indicies = findQueueFamilies(device, mVulkanSurface);

		bool extensionsSupported = checkDeviceExtensionSupport(device);

		int score = 0;

		if (deviceProperties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU) score += 1000;
		else if (deviceProperties.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU) score += 500;

		score += deviceProperties.limits.maxImageDimension2D;

		if (!deviceFeatures.geometryShader) return 0;

		if (!indicies.isComplete()) return 0;

		if (!extensionsSupported) return 0;

		bool swapChainAdequate = false;

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

		if (!swapChainAdequate) return 0;

		return score;
	}

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

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

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

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

		return requiredExtensions.empty();
	}

	VkResult Aethiopicus::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 Aethiopicus::destroyDebugUtilsMessengerEXT(VkInstance instance, VkDebugUtilsMessengerEXT debugMessenger, const VkAllocationCallbacks* pAllocator) {
		auto func = (PFN_vkDestroyDebugUtilsMessengerEXT) vkGetInstanceProcAddr(instance, "vkDestroyDebugUtilsMessengerEXT");
		if (func != nullptr) {
			func(instance, debugMessenger, pAllocator);
		}
	}

	void Aethiopicus::clearScreen(Color col) {}

	void Aethiopicus::renderCurrent() {}

	void Aethiopicus::drawText(Resource* resource, std::string text, IntVec2 pos, Color color) {}
	void Aethiopicus::drawTexture(Resource* resource, IntRect size) {}

	void Aethiopicus::useShader(Shader* shader, Point2D points[3]) {}
	void Aethiopicus::stopUsingShaders() {}

	// Pre and Post draw
	void Aethiopicus::preDraw() {}
	void Aethiopicus::postDraw() {}

	// Update
	void Aethiopicus::update() {}

	// Util
	void Aethiopicus::compileShader(Shader* shader) {}
	void Aethiopicus::toggleWireframe() {}
	void Aethiopicus::setWindowTitle(std::string title) {}
	void Aethiopicus::resize(unsigned int width, unsigned int height) {}
}