#include "Renderer.hpp"

#include "UI.hpp"
#include "../../Window.hpp"
#include "shared.hpp"
#include "SwapChain.hpp"
#include "Pipeline.hpp"
#include "CommandPool.hpp"
#include <GLFW/glfw3.h>
#include <string.h>
#include <algorithm>
#include <set>

using namespace render::vk;

constexpr auto LOAD_DEVICE = true;
constexpr auto VALIDATION_LAYER = true;

void set_current_extent(VkSurfaceCapabilitiesKHR &capabilities, GLFWwindow *ptr);

Renderer::Renderer(VkInstance instance, VkDevice device, const PhysicalDeviceInfo& info, const renderOptions& opt):
        options(opt), instance(instance), surface(info.surface), device(device),
        physicalInfo(std::make_unique<PhysicalDeviceInfo>(info)) {
    vkGetDeviceQueue(device, info.queueIndices.graphicsFamily.value(), 0, &graphicsQueue);
    vkGetDeviceQueue(device, info.queueIndices.presentFamily.value(), 0, &presentQueue);

    set_current_extent(physicalInfo->swapDetails.capabilities, physicalInfo->window);
    physicalInfo->surfaceFormat = [&]() {
        for(const auto& format: info.swapDetails.formats) {
            if (format.format == VK_FORMAT_B8G8R8A8_SRGB && format.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
                return format;
            }
        }

        LOG_W("Using suboptimal surface format");
        return info.swapDetails.formats[0];
    }();

    swapChain = std::make_unique<SwapChain>(device, *physicalInfo.get());
    pipeline = std::make_unique<Pipeline>(device, *physicalInfo.get(), options);
    commandPool = std::make_unique<CommandPool>(device, swapChain->getImageViews(), pipeline->getRef(), *physicalInfo.get(), options);

    {
        imageAvailableSemaphores.resize(opt.inFlightFrames);
        renderFinishedSemaphores.resize(opt.inFlightFrames);
        inFlightFences.resize(opt.inFlightFrames);

        VkSemaphoreCreateInfo semaphoreInfo{};
        semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

        VkFenceCreateInfo fenceInfo{};
        fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
        fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

        for (int i = 0; i < opt.inFlightFrames; i++) {
            if (vkCreateSemaphore(device, &semaphoreInfo, ALLOC, &imageAvailableSemaphores[i]) != VK_SUCCESS ||
                vkCreateSemaphore(device, &semaphoreInfo, ALLOC, &renderFinishedSemaphores[i]) != VK_SUCCESS ||
                vkCreateFence(device, &fenceInfo, ALLOC, &inFlightFences[i]) != VK_SUCCESS)
            {
                FATAL("Failed to create synchronization objects!");
            }
        }
        vkResetFences(device, 1, &inFlightFences[currentFrame]);
    }
}
Renderer::~Renderer() {
    vkDeviceWaitIdle(device);
    destroySwapChain();

    for(size_t i = 0; i < renderFinishedSemaphores.size(); i++) {
        vkDestroyFence(device, inFlightFences[i], ALLOC);
        vkDestroySemaphore(device, renderFinishedSemaphores[i], ALLOC);
        vkDestroySemaphore(device, imageAvailableSemaphores[i], ALLOC);
    }

    commandPool.reset();

    vkDestroyDevice(device, ALLOC);
    vkDestroySurfaceKHR(instance, surface, ALLOC);
    vkDestroyInstance(instance, ALLOC);
}

void Renderer::recreateSwapChain() {
    LOG_D("Recreating swapchain");
    vkDeviceWaitIdle(device);
    destroySwapChain();

    physicalInfo->swapDetails = SwapChainSupportDetails::Query(physicalInfo->device, physicalInfo->surface);
    set_current_extent(physicalInfo->swapDetails.capabilities, physicalInfo->window);
    swapChain = std::make_unique<SwapChain>(device, *physicalInfo.get());
    pipeline = std::make_unique<Pipeline>(device, *physicalInfo.get(), options);
    commandPool->allocate(swapChain->getImageViews(), pipeline->getRef(), physicalInfo->swapDetails.capabilities.currentExtent, options);
}
void Renderer::destroySwapChain() {
    commandPool->free();
    pipeline.reset();
    swapChain.reset();
}

void on_resize_callback(GLFWwindow *, int, int) {
    Renderer::Get()->setResized();
}
void set_current_extent(VkSurfaceCapabilitiesKHR &capabilities, GLFWwindow* ptr) {
    if(capabilities.currentExtent.width != INT32_MAX) {
        return;
    }
    auto windowSize = std::make_pair<int, int>(0, 0);
    glfwGetFramebufferSize(ptr, &windowSize.first, &windowSize.second);
    capabilities.currentExtent = VkExtent2D{
        std::max(capabilities.minImageExtent.width, std::min<uint32_t>(capabilities.maxImageExtent.width, windowSize.first)),
        std::max(capabilities.minImageExtent.height, std::min<uint32_t>(capabilities.maxImageExtent.height, windowSize.second))};
};

bool Renderer::Load(Window& window, const renderOptions& opt) {
    Window::CreateInfo windowInfo;
    windowInfo.pfnResize = on_resize_callback;
    windowInfo.client = {Window::CreateInfo::Client::Type::VK, 0, 0};
    windowInfo.samples = 1;
    if (!window.create(windowInfo))
        return false;

    if (volkInitialize() != VK_SUCCESS) {
        LOG_E("Failed to initialize Vulkan");
        return false;
    }

    VkInstance instance;
    std::vector<const char *> layers;
    {
        VkApplicationInfo appInfo{};
        appInfo.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
        appInfo.pApplicationName = "Univerxel";
        appInfo.applicationVersion = VK_MAKE_VERSION(0, 0, 1);
        appInfo.pEngineName = "No Engine";
        appInfo.engineVersion = VK_MAKE_VERSION(1, 0, 0);
        appInfo.apiVersion = VK_API_VERSION_1_2;

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

        std::vector<const char *> extensions;
        { // Check extensions
            uint32_t availableExtensionCount = 0;
            vkEnumerateInstanceExtensionProperties(nullptr, &availableExtensionCount, nullptr);
            std::vector<VkExtensionProperties> availableExtensions(availableExtensionCount);
            vkEnumerateInstanceExtensionProperties(nullptr, &availableExtensionCount, availableExtensions.data());

#if DEBUG
            LOG_D("Available extensions:");
            for (const auto &extension : availableExtensions) {
                LOG_D('\t' << extension.extensionName << " : " << extension.specVersion);
            }
#endif

            uint32_t glfwExtensionCount = 0;
            const char **glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);
            extensions.reserve(glfwExtensionCount);
            for (uint32_t i = 0; i < glfwExtensionCount; i++) {
                if (std::none_of(availableExtensions.begin(), availableExtensions.end(), [&](const VkExtensionProperties &ex) { return strcmp(ex.extensionName, glfwExtensions[i]) == 0; })) {
                    LOG_E("Missing required glfw extension " << glfwExtensions[i]);
                    return false;
                }
                extensions.push_back(glfwExtensions[i]);
            }

        }
        createInfo.enabledExtensionCount = extensions.size();
        createInfo.ppEnabledExtensionNames = extensions.data();

        { // Check layers
            uint32_t availableLayerCount = 0;
            vkEnumerateInstanceLayerProperties(&availableLayerCount, nullptr);
            std::vector<VkLayerProperties> availableLayers(availableLayerCount);
            vkEnumerateInstanceLayerProperties(&availableLayerCount, availableLayers.data());

#if DEBUG
            LOG_D("Available layers:");
            for (const auto &layer : availableLayers) {
                LOG_D('\t' << layer.layerName << " : " << layer.specVersion);
            }
#endif

            const auto hasLayer = [&availableLayers](const char *layer) {
                return std::any_of(availableLayers.begin(), availableLayers.end(), [&layer](const VkLayerProperties &l) { return strcmp(l.layerName, layer) == 0; });
            };

            if constexpr (VALIDATION_LAYER) {
                constexpr auto VALIDATION_LAYER_NAME = "VK_LAYER_KHRONOS_validation";
                if (hasLayer(VALIDATION_LAYER_NAME)) {
                    layers.push_back(VALIDATION_LAYER_NAME);
                } else {
                    LOG_W("Validation layer unavailable");
                }
                if (hasLayer(VK_EXT_DEBUG_UTILS_EXTENSION_NAME)) {
                    layers.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
                } else {
                    LOG_W("Debug utils layer unavailable");
                }
            }
        }
        createInfo.enabledLayerCount = layers.size();
        createInfo.ppEnabledLayerNames = layers.data();

        if (vkCreateInstance(&createInfo, ALLOC, &instance) != VK_SUCCESS) {
            LOG_E("Failed to create Vulkan instance");
            return false;
        }
    }

    if constexpr(LOAD_DEVICE) {
        volkLoadInstanceOnly(instance);
    } else {
        volkLoadInstance(instance);
    }

    const auto version = volkGetInstanceVersion();
    LOG_D("Vulkan " << VK_VERSION_MAJOR(version) << '.' << VK_VERSION_MINOR(version) << '.' << VK_VERSION_PATCH(version) << ", GLSL TODO:");

    VkSurfaceKHR surface;
    if (glfwCreateWindowSurface(instance, window.getPtr(), ALLOC, &surface) != VK_SUCCESS) {
        LOG_E("Failed to create window surface!");
        return false;
    }

    PhysicalDeviceInfo physicalInfo;
    std::vector<const char *> requiredExtensions = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};
    {
        uint32_t deviceCount = 0;
        vkEnumeratePhysicalDevices(instance, &deviceCount, nullptr);
        if (deviceCount == 0) {
            LOG_E("Any GPU with Vulkan support");
            return false;
        }
        std::vector<VkPhysicalDevice> devices(deviceCount);
        vkEnumeratePhysicalDevices(instance, &deviceCount, devices.data());

        uint bestScore = 0;
        for(const auto& device: devices) {
            uint score = 1;
            VkPhysicalDeviceProperties deviceProperties;
            VkPhysicalDeviceFeatures deviceFeatures;
            vkGetPhysicalDeviceProperties(device, &deviceProperties);
            vkGetPhysicalDeviceFeatures(device, &deviceFeatures);

            if (!deviceFeatures.geometryShader)
                continue;

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

                if (std::any_of(requiredExtensions.begin(), requiredExtensions.end(), [&](const char *required) {
                    return std::none_of(availableExtensions.begin(), availableExtensions.end(), [&](const VkExtensionProperties &ex) {
                        return strcmp(ex.extensionName, required) == 0;
                    });
                }))
                    continue;
            }

            if (deviceProperties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
                score += 10000;

            score += deviceProperties.limits.maxImageDimension2D;
            //TODO: check others limits

            auto infos = PhysicalDeviceInfo{device, window.getPtr(),
                SwapChainSupportDetails::Query(device, surface), QueueFamilyIndices::Query(device, surface),
                VkSurfaceFormatKHR(), surface};

            if (!infos.queueIndices.isComplete())
                continue;

            if (!infos.swapDetails.isValid())
                continue;

            if (score > bestScore) {
                bestScore = score;
                physicalInfo = infos;
            }
        }
        if (physicalInfo.device == VK_NULL_HANDLE) {
            LOG_E("Any GPU matching requirements");
            return false;
        }
        VkPhysicalDeviceProperties deviceProperties;
        vkGetPhysicalDeviceProperties(physicalInfo.device, &deviceProperties);
        LOG_D("Using " << deviceProperties.deviceName);
    }

    VkDevice device;
    {
        std::vector<VkDeviceQueueCreateInfo> queueCreateInfos;
        {
            std::set<uint32_t> uniqueQueueFamilies = {physicalInfo.queueIndices.graphicsFamily.value(), physicalInfo.queueIndices.presentFamily.value()};

            const 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{};
        //TODO:

        VkDeviceCreateInfo createInfo{};
        createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
        createInfo.pQueueCreateInfos = queueCreateInfos.data();
        createInfo.queueCreateInfoCount = queueCreateInfos.size();
        createInfo.pEnabledFeatures = &deviceFeatures;
        createInfo.enabledLayerCount = layers.size();
        createInfo.ppEnabledLayerNames = layers.data();
        createInfo.enabledExtensionCount = requiredExtensions.size();
        createInfo.ppEnabledExtensionNames = requiredExtensions.data();

        if (vkCreateDevice(physicalInfo.device, &createInfo, ALLOC, &device) != VK_SUCCESS) {
            LOG_E("Failed to bind graphic device");
            return false;
        }
    }

    if constexpr(LOAD_DEVICE) {
        volkLoadDevice(device);
    }

    sInstance = new Renderer(instance, device, physicalInfo, opt);
    return true;
}

void Renderer::loadUI(Window& w) {
    UI::Load(w);
}

SwapChainSupportDetails SwapChainSupportDetails::Query(VkPhysicalDevice device, VkSurfaceKHR surface) {
    SwapChainSupportDetails swapDetails;
    vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, surface, &swapDetails.capabilities);

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

    uint32_t presentModeCount;
    vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, nullptr);
    if (presentModeCount != 0) {
        swapDetails.presentModes.resize(presentModeCount);
        vkGetPhysicalDeviceSurfacePresentModesKHR(device, surface, &presentModeCount, swapDetails.presentModes.data());
    }
    return swapDetails;
}

QueueFamilyIndices QueueFamilyIndices::Query(VkPhysicalDevice device, VkSurfaceKHR surface) {
    QueueFamilyIndices queueIndices;
    uint32_t queueFamilyCount = 0;
    vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, nullptr);
    std::vector<VkQueueFamilyProperties> queueFamilies(queueFamilyCount);
    vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies.data());

    int i = 0;
    for(const auto& queueFamily: queueFamilies) {
        if (queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT) {
            queueIndices.graphicsFamily = i;
        }
        VkBool32 presentSupport = false;
        vkGetPhysicalDeviceSurfaceSupportKHR(device, i, surface, &presentSupport);
        if (presentSupport) {
            queueIndices.presentFamily = i;
        }

#if DEBUG
        LOG_D("Queue " << i << ' ' << (queueFamily.queueFlags & VK_QUEUE_GRAPHICS_BIT ? "graphics " : "")
                    << (queueFamily.queueFlags & VK_QUEUE_COMPUTE_BIT ? "compute " : "")
                    << (queueFamily.queueFlags & VK_QUEUE_TRANSFER_BIT ? "transfer " : "")
                    << (queueFamily.queueFlags & VK_QUEUE_PROTECTED_BIT ? "protected " : "")
                    << 'x' << queueFamily.queueCount);
#endif
        i++;
    }

    return queueIndices;
}

void Renderer::beginFrame() {
    assert(currentImage == UINT32_MAX);

    //FIXME: TracyGpuZone("Render");
    if (auto newImage = swapChain->acquireNextImage(imageAvailableSemaphores[currentFrame], inFlightFences[currentFrame])) {
        currentImage = newImage.value();
    } else {
        recreateSwapChain();
        beginFrame();
    }
}

std::function<buffer::params(glm::mat4)> Renderer::beginWorldPass() {
    assert(currentImage < swapChain->getImageViews().size());

    commandPool->submitGraphics(currentImage, graphicsQueue, imageAvailableSemaphores[currentFrame],
        renderFinishedSemaphores[currentFrame], inFlightFences[currentFrame]);

    /*WorldPass->useIt();
    WorldPass->start(this);*/
    return [&](glm::mat4) {
        return buffer::params{}; //WorldPass->setup(this, model);
    };
}

std::function<buffer::params(const std::vector<glm::mat4> &)> Renderer::beginEntityPass() {
    /*EntityPass->useIt();
    EntityPass->start(this);*/
    return [&](const std::vector<glm::mat4>&) {
        return buffer::params{}; //EntityPass->setup(this, models);
    };
}

size_t Renderer::drawIndicatorCube(glm::mat4) {
    /*IndicatorPass->useIt();
    return IndicatorCubeBuffer.draw(IndicatorPass->setup(this, model));*/
    return 0;
}

void Renderer::endPass() {
    /*if(SkyEnable) {
        SkyPass->draw(this);
    }*/
}

void Renderer::swapBuffer(Window&) {
    /*TracyGpuZone("Swap");
    glfwSwapBuffers(ptr);
    TracyGpuCollect;*/

    if(!swapChain->presentImage(currentImage, presentQueue, renderFinishedSemaphores[currentFrame]) || framebufferResized) {
        framebufferResized = false;
        recreateSwapChain();
    }

    currentFrame = (currentFrame + 1) % renderFinishedSemaphores.size();
    currentImage = UINT32_MAX;

    vkWaitForFences(device, 1, &inFlightFences[currentFrame], VK_TRUE, UINT64_MAX);
}

void Renderer::reloadShaders(const passOptions&) {
    /*WorldPass = std::make_unique<pass::WorldProgram>(options);
    EntityPass = std::make_unique<pass::EntityProgram>(options);*/
}
void Renderer::reloadTextures(const std::string&, float, float) {
    /*unloadTextures();
    loadTextures(texturePath, mipMapLOD, anisotropy);*/
}

void Renderer::lookFrom(const Camera&) {
    /*ProjectionMatrix = camera.getProjectionMatrix();
    ViewMatrix = camera.getViewMatrix();
    FogDepth = camera.getDepth();*/
}

void Renderer::setClearColor(glm::vec4) {
    /*FogColor = c;
    glClearColor(c.r, c.g, c.b, c.a);*/
}
void Renderer::setCurvature(glm::vec4, float) {
    /*SphereProj = sp;
    Curvature = c;*/
}