#include "renderer.h"

#include <set>
#include <vector>
#include <iostream>
#include <cstring>

#include "platform.h"

Renderer::Renderer() {
    createInstance();
#ifdef DEBUG
    if(enableDebug)
        createDebugMessenger();
#endif
    createLogicalDevice();
    createCommandPool();
}

Renderer::~Renderer() {
#ifdef DEBUG
    destroyMessenger_(instance_, messenger_, nullptr);
#endif
    
    vkDestroyDevice(device_, nullptr);
    
    vkDestroyInstance(instance_, nullptr);
}

void Renderer::render(RenderTarget* target) {
    uint32_t imageIndex = 0;
    vkAcquireNextImageKHR(device_, target->swapchain, UINT64_MAX, target->imageAvailableSemaphore, nullptr, &imageIndex);

    vkWaitForFences(device_, 1, &target->fences[imageIndex], true, UINT64_MAX);
    vkResetFences(device_, 1, &target->fences[imageIndex]);
    
    const VkCommandBuffer commandBuffer = target->commandBuffers[imageIndex];

    VkCommandBufferBeginInfo beginInfo = {};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;

    vkBeginCommandBuffer(commandBuffer, &beginInfo);

    vkEndCommandBuffer(commandBuffer);

    const VkPipelineStageFlags waitStage = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    
    VkSubmitInfo submitInfo = {};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.waitSemaphoreCount = 1;
    submitInfo.pWaitSemaphores = &target->imageAvailableSemaphore;
    submitInfo.pWaitDstStageMask = &waitStage;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &commandBuffer;
    submitInfo.signalSemaphoreCount = 1;
    submitInfo.pSignalSemaphores = &target->renderFinishedSemaphore;

    vkQueueSubmit(graphicsQueue_, 1, &submitInfo, target->fences[imageIndex]);
    
    VkPresentInfoKHR presentInfo = {};
    presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
    presentInfo.waitSemaphoreCount = 1;
    presentInfo.pWaitSemaphores = &target->renderFinishedSemaphore;
    presentInfo.swapchainCount = 1;
    presentInfo.pSwapchains = &target->swapchain;
    presentInfo.pImageIndices = &imageIndex;

    vkQueuePresentKHR(graphicsQueue_, &presentInfo);
}

RenderTarget* Renderer::createSurfaceRenderTarget(VkSurfaceKHR surface) {
    RenderTarget* target = new RenderTarget();
    target->surface = surface;
    
    VkSurfaceCapabilitiesKHR surfaceCapabilities = {};
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physicalDevice_, surface, &surfaceCapabilities);
    
    uint32_t surfaceFormatCount = 0;
    vkGetPhysicalDeviceSurfaceFormatsKHR(physicalDevice_, surface, &surfaceFormatCount, nullptr);
    
    VkSurfaceFormatKHR* surfaceFormats = new VkSurfaceFormatKHR[surfaceFormatCount];
    vkGetPhysicalDeviceSurfaceFormatsKHR(physicalDevice_, surface, &surfaceFormatCount, surfaceFormats);
    
    uint32_t chosenFormat = 0;
    for(uint32_t i = 0; i < surfaceFormatCount; i++) {
        if(surfaceFormats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR)
            chosenFormat = i;
    }
    
    VkBool32 supported = false;
    vkGetPhysicalDeviceSurfaceSupportKHR(physicalDevice_, queueIndices.presentation, surface, &supported);
    
    VkSwapchainCreateInfoKHR swapchainCreateInfo = {};
    swapchainCreateInfo.sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR;
    swapchainCreateInfo.surface = surface;
    swapchainCreateInfo.minImageCount = surfaceCapabilities.minImageCount;
    swapchainCreateInfo.imageColorSpace = surfaceFormats[chosenFormat].colorSpace;
    swapchainCreateInfo.imageFormat = surfaceFormats[chosenFormat].format;
    swapchainCreateInfo.imageExtent = surfaceCapabilities.currentExtent;
    swapchainCreateInfo.imageArrayLayers = 1;
    swapchainCreateInfo.imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
    swapchainCreateInfo.queueFamilyIndexCount = 1;
    swapchainCreateInfo.pQueueFamilyIndices = &queueIndices.presentation;
    swapchainCreateInfo.preTransform = surfaceCapabilities.currentTransform;
    swapchainCreateInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
    swapchainCreateInfo.presentMode = VK_PRESENT_MODE_FIFO_KHR;
    swapchainCreateInfo.clipped = true;
    
    vkCreateSwapchainKHR(device_, &swapchainCreateInfo, nullptr, &target->swapchain);
    
    delete[] surfaceFormats;
    
    uint32_t swapchainImageCount = 0;
    vkGetSwapchainImagesKHR(device_, target->swapchain, &swapchainImageCount, nullptr);
    
    VkImage* swapchainImages = new VkImage[swapchainImageCount];
    vkGetSwapchainImagesKHR(device_, target->swapchain, &swapchainImageCount, swapchainImages);
    
    VkCommandBufferAllocateInfo allocateInfo = {};
    allocateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocateInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocateInfo.commandPool = commandPool_;
    allocateInfo.commandBufferCount = swapchainImageCount;
    
    target->commandBuffers = new VkCommandBuffer[swapchainImageCount];
    
    vkAllocateCommandBuffers(device_, &allocateInfo, target->commandBuffers);
    
    VkSemaphoreCreateInfo semaphoreCreateInfo = {};
    semaphoreCreateInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;
    
    vkCreateSemaphore(device_, &semaphoreCreateInfo, nullptr, &target->imageAvailableSemaphore);
    vkCreateSemaphore(device_, &semaphoreCreateInfo, nullptr, &target->renderFinishedSemaphore);
    
    VkFenceCreateInfo fenceCreateInfo = {};
    fenceCreateInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
    fenceCreateInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

    target->fences = new VkFence[swapchainImageCount];
    
    for(uint32_t i = 0; i < swapchainImageCount; i++)
        vkCreateFence(device_, &fenceCreateInfo, nullptr, &target->fences[i]);

    return target;
}

void Renderer::destroyRenderTarget(RenderTarget* target) {
    vkDestroySwapchainKHR(device_, target->swapchain, nullptr);
    vkDestroySurfaceKHR(instance_, target->surface, nullptr);
    
    delete target;
}

void Renderer::createInstance() {
    uint32_t layerCount = 0;
    vkEnumerateInstanceLayerProperties(&layerCount, nullptr);
    
    VkLayerProperties* availableLayers = new VkLayerProperties[layerCount];
    vkEnumerateInstanceLayerProperties(&layerCount, availableLayers);
    
    std::vector<const char*> enabledLayers;
#ifdef DEBUG
    for(uint32_t i = 0; i < layerCount; i++) {
        if(!strcmp(availableLayers[i].layerName, "VK_LAYER_LUNARG_standard_validation"))
            enabledLayers.push_back("VK_LAYER_LUNARG_standard_validation");
    }
#endif

    delete[] availableLayers;
    
    uint32_t extensionCount = 0;
    vkEnumerateInstanceExtensionProperties(nullptr, &extensionCount, nullptr);
    
    VkExtensionProperties* availableExtensions = new VkExtensionProperties[extensionCount];
    vkEnumerateInstanceExtensionProperties(nullptr, &extensionCount, availableExtensions);

    std::vector<const char*> enabledExtensions;
#ifdef DEBUG
    for(uint32_t i = 0; i < extensionCount; i++) {
        if(!strcmp(availableExtensions[i].extensionName, "VK_EXT_debug_utils")) {
            enabledExtensions.push_back("VK_EXT_debug_utils");
            enableDebug = true;
        }
    }
#endif

    delete[] availableExtensions;
    
    auto requiredExtensions = platform::getRequiredExtensions();
    enabledExtensions.insert(enabledExtensions.end(), requiredExtensions.begin(), requiredExtensions.end());

    VkInstanceCreateInfo instanceCreateInfo = {};
    instanceCreateInfo.sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO;
    instanceCreateInfo.enabledLayerCount = enabledLayers.size();
    instanceCreateInfo.ppEnabledLayerNames = enabledLayers.data();
    instanceCreateInfo.enabledExtensionCount = enabledExtensions.size();
    instanceCreateInfo.ppEnabledExtensionNames = enabledExtensions.data();
    
    vkCreateInstance(&instanceCreateInfo, nullptr, &instance_);
}

#ifdef DEBUG
void Renderer::createDebugMessenger() {
    createMessenger_ = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance_, "vkCreateDebugUtilsMessengerEXT");
    destroyMessenger_ = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(instance_, "vkDestroyDebugUtilsMessengerEXT");

    VkDebugUtilsMessengerCreateInfoEXT messengerCreateInfo = {};
    messengerCreateInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
    messengerCreateInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT;
    messengerCreateInfo.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;
    messengerCreateInfo.pfnUserCallback = [](VkDebugUtilsMessageSeverityFlagBitsEXT, unsigned int, const VkDebugUtilsMessengerCallbackDataEXT* callback, void*) -> unsigned int {
        std::cout << callback->pMessage << std::endl;
        
        return VK_SUCCESS;
    };
    
    createMessenger_(instance_, &messengerCreateInfo, nullptr, &messenger_);
}
#endif

void Renderer::createLogicalDevice() {
    uint32_t physicalDeviceCount = 0;
    vkEnumeratePhysicalDevices(instance_, &physicalDeviceCount, nullptr);
    
    VkPhysicalDevice* physicalDevices = new VkPhysicalDevice[physicalDeviceCount];
    vkEnumeratePhysicalDevices(instance_, &physicalDeviceCount, physicalDevices);
        
    for(uint32_t i = 0; i < physicalDeviceCount; i++) {
        VkPhysicalDeviceProperties properties = {};
        vkGetPhysicalDeviceProperties(physicalDevices[i], &properties);
    
        physicalDevice_ = physicalDevices[i];
    }
    
    delete[] physicalDevices;
    
    uint32_t queueFamilyPropertiesCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice_, &queueFamilyPropertiesCount, nullptr);
    
    VkQueueFamilyProperties* queueFamilyProperties = new VkQueueFamilyProperties[queueFamilyPropertiesCount];
    vkGetPhysicalDeviceQueueFamilyProperties(physicalDevice_, &queueFamilyPropertiesCount, queueFamilyProperties);
        
    for(uint32_t i = 0; i < queueFamilyPropertiesCount; i++) {
        if(queueFamilyProperties[i].queueFlags & VK_QUEUE_GRAPHICS_BIT)
            queueIndices.graphics = i;
    }
    
    delete[] queueFamilyProperties;
    
    queueIndices.presentation = queueIndices.graphics; //FIXME: this may not always be true!!
    
    const std::set<uint32_t> queueFamilyIndices = {queueIndices.graphics};
    std::vector<VkDeviceQueueCreateInfo> deviceQueueCreateInfos;
    for(auto queueFamilyIndex : queueFamilyIndices) {
        const float priority = 1.0f;
        
        VkDeviceQueueCreateInfo queueCreateInfo = {};
        queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
        queueCreateInfo.queueFamilyIndex = queueFamilyIndex;
        queueCreateInfo.queueCount = 1;
        queueCreateInfo.pQueuePriorities = &priority;
        
        deviceQueueCreateInfos.push_back(queueCreateInfo);
    }
    
    const std::vector<const char*> enabledExtensions = {"VK_KHR_swapchain"};
    
    VkDeviceCreateInfo deviceCreateInfo = {};
    deviceCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    deviceCreateInfo.queueCreateInfoCount = deviceQueueCreateInfos.size();
    deviceCreateInfo.pQueueCreateInfos = deviceQueueCreateInfos.data();
    deviceCreateInfo.enabledExtensionCount = enabledExtensions.size();
    deviceCreateInfo.ppEnabledExtensionNames = enabledExtensions.data();
    
    vkCreateDevice(physicalDevice_, &deviceCreateInfo, nullptr, &device_);
    
    vkGetDeviceQueue(device_, queueIndices.graphics, 0, &graphicsQueue_);
}

void Renderer::createCommandPool() {
    VkCommandPoolCreateInfo poolCreateInfo = {};
    poolCreateInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolCreateInfo.queueFamilyIndex = queueIndices.graphics;
    
    vkCreateCommandPool(device_, &poolCreateInfo, nullptr, &commandPool_);
}