#include "Renderer.hpp"

#include "UI.hpp"
#include "../../Window.hpp"
#include "../../control/Camera.hpp"
#include "PhysicalDeviceInfo.hpp"
#include "Allocator.hpp"
#include "SwapChain.hpp"
#include "Pipeline.hpp"
#include "CommandCenter.hpp"
#include <GLFW/glfw3.h>
#include <string.h>
#include <algorithm>
#include <set>
#include <Tracy.hpp>

using namespace render::vk;

constexpr auto LOAD_DEVICE = true;
#if LOG_DEBUG
constexpr auto VALIDATION_LAYER = true;
#else
constexpr auto VALIDATION_LAYER = false;
#endif

void set_current_extent(VkSurfaceCapabilitiesKHR &capabilities, GLFWwindow *ptr);
VKAPI_ATTR VkBool32 VKAPI_CALL debugValidationCallback(
    VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
    VkDebugUtilsMessageTypeFlagsEXT messageType,
    const VkDebugUtilsMessengerCallbackDataEXT *pCallbackData,
    void *pUserData);

Renderer::Renderer(VkInstance instance, VkDevice device, const PhysicalDeviceInfo& info, const render::renderOptions& opt):
        options(opt), instance(instance), surface(info.surface), device(device),
        physicalInfo(std::make_unique<PhysicalDeviceInfo>(info)) {
    if constexpr(VALIDATION_LAYER) {
        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_INFO_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 = debugValidationCallback;
        createInfo.pUserData = nullptr;

        if (vkCreateDebugUtilsMessengerEXT(instance, &createInfo, ALLOC, &debugMessenger) != VK_SUCCESS) {
            LOG_E("Failed to redirect validation errors");
        }
    }

    set_current_extent(physicalInfo->swapDetails.capabilities, physicalInfo->window);

    allocator = std::make_unique<Allocator>(device, getInfos());
    Allocator::MakeDefault(allocator.get());
    swapChain = std::make_unique<SwapChain>(device, getInfos());
    pipeline = std::make_unique<Pipeline>(device, getInfos(), options);
    commandCenter = std::make_unique<CommandCenter>(device, getInfos(), options);
    commandCenter->allocate(swapChain->getImageViews(), *pipeline.get(), getInfos().device, getInfos().swapDetails.capabilities.currentExtent);

    {
        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);
    const auto imageCount = swapChain->getImageViews().size();
    destroySwapChain();
    render::UI::Unload();

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

    commandCenter.reset();
    for (size_t i = 0; i < imageCount; i++) {
        ShortIndexedVertexBuffer::ClearUnused(i);
    }
    allocator.reset();

    vkDestroyDevice(device, ALLOC);
    vkDestroySurfaceKHR(instance, surface, ALLOC);
    if constexpr(VALIDATION_LAYER) {
        vkDestroyDebugUtilsMessengerEXT(instance, debugMessenger, ALLOC);
    }
    vkDestroyInstance(instance, ALLOC);
}

void Renderer::recreateSwapChain() {
    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, getInfos());
    pipeline = std::make_unique<Pipeline>(device, getInfos(), options);
    commandCenter->allocate(swapChain->getImageViews(), *pipeline.get(), getInfos().device, getInfos().swapDetails.capabilities.currentExtent);

    dynamic_cast<UI*>(UI::Get())->setImageCount(swapChain->getImageViews().size());
}
void Renderer::destroySwapChain() {
    commandCenter->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))};
};
VKAPI_ATTR VkBool32 VKAPI_CALL debugValidationCallback(VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity, VkDebugUtilsMessageTypeFlagsEXT,
    const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData, void*)
{
    switch (messageSeverity) {
    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT:
        LOG_E("[VK] " << pCallbackData->pMessage);
        break;

    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT:
        LOG_W("[VK] " << pCallbackData->pMessage);
        break;

    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_INFO_BIT_EXT:
        LOG_D("[VK] " << pCallbackData->pMessage);
        break;

    case VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT:
        LOG_T("[VK] " << pCallbackData->pMessage);
        break;

    default:
        LOG_I("[VK] " << pCallbackData->pMessage);
        break;
    }

    return VK_FALSE;
}

bool Renderer::Load(Window& window, const render::renderOptions& opt, const windowOptions& windOpt) {
    Window::CreateInfo windowInfo;
    windowInfo.pfnResize = on_resize_callback;
    windowInfo.client = {Window::CreateInfo::Client::Type::VK, 0, 0};
    windowInfo.samples = -1;
    windowInfo.fps = windOpt.targetFPS;
    windowInfo.fullscreen = windOpt.fullscreen;

    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_MAKE_VERSION(1, 2, 0);

        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 LOG_TRACE
            LOG_T("Available instance extensions:");
            for (const auto &extension : availableExtensions) {
                LOG_T('\t' << extension.extensionName << " : " << extension.specVersion);
            }
#endif

            const auto hasExtension = [&availableExtensions](const char *extension) {
                return std::any_of(availableExtensions.begin(), availableExtensions.end(), [&extension](const VkExtensionProperties &ex) { return strcmp(ex.extensionName, extension) == 0; });
            };

            uint32_t glfwExtensionCount = 0;
            const char **glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);
            extensions.reserve(glfwExtensionCount);
            for (uint32_t i = 0; i < glfwExtensionCount; i++) {
                if (!hasExtension(glfwExtensions[i])) {
                    LOG_E("Missing required glfw extension " << glfwExtensions[i]);
                    return false;
                }
                extensions.push_back(glfwExtensions[i]);
            }
            if constexpr (VALIDATION_LAYER) {
                if (hasExtension(VK_EXT_DEBUG_UTILS_EXTENSION_NAME)) {
                    extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
                } else {
                    LOG_W("Debug utils extension unavailable");
                }
            }
        }
        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 LOG_TRACE
            LOG_T("Available layers:");
            for (const auto &layer : availableLayers) {
                LOG_T('\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");
                }
            }
        }
        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);
    }

    [[maybe_unused]]
    const auto version = volkGetInstanceVersion();
    LOG_D("Vulkan " << VK_VERSION_MAJOR(version) << '.' << VK_VERSION_MINOR(version) << '.' << VK_VERSION_PATCH(version) << ", GLSL precompiled");

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

    PhysicalDeviceInfo physicalInfo;
    const std::vector<const char *> requiredExtensions = {VK_KHR_SWAPCHAIN_EXTENSION_NAME};
    const std::vector<const char *> optionalExtensions = {VK_KHR_GET_MEMORY_REQUIREMENTS_2_EXTENSION_NAME, VK_EXT_MEMORY_BUDGET_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());

        uint32_t bestScore = 0;
        for(const auto& device: devices) {
            uint32_t score = 1;
            auto infos = PhysicalDeviceInfo(window.getPtr(), device, surface, windOpt.getSamples(), windOpt.targetFPS < Window::MIN_FPS);

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

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

                const auto hasExtension = [&availableExtensions](const char *extension) {
                    return std::any_of(availableExtensions.begin(), availableExtensions.end(), [&extension](const VkExtensionProperties &ex) { return strcmp(ex.extensionName, extension) == 0; });
                };

                if (std::any_of(requiredExtensions.begin(), requiredExtensions.end(), [&](const char *required) { return !hasExtension(required); }))
                    continue;

                for (auto extension: optionalExtensions) {
                    if (hasExtension(extension))
                        infos.optionalExtensions.push_back(extension);
                }
            }

            if (infos.properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU)
                score += 10000;

            score += infos.properties.limits.maxImageDimension2D;
            score += infos.samples * 50;

            if (infos.features.geometryShader)
                score += 2000;

            if (infos.features.samplerAnisotropy)
                score += 1000;

            if (infos.features.sampleRateShading)
                score += 1500;

            if (infos.features.fillModeNonSolid)
                score += 100;

            //TODO: check others limits

            if (!infos.queueIndices.isComplete())
                continue;
            if (infos.queueIndices.isOptimal())
                score += 5000;

            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;
        }
        LOG_D("Using " << physicalInfo.properties.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{};
        deviceFeatures.geometryShader = physicalInfo.features.geometryShader;
        deviceFeatures.samplerAnisotropy = physicalInfo.features.samplerAnisotropy;
        deviceFeatures.sampleRateShading = physicalInfo.features.sampleRateShading;
        deviceFeatures.fillModeNonSolid = physicalInfo.features.fillModeNonSolid;

        std::vector<const char*> extensions(requiredExtensions);
        extensions.insert(extensions.end(), physicalInfo.optionalExtensions.begin(), physicalInfo.optionalExtensions.end());
        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 = extensions.size();
        createInfo.ppEnabledExtensionNames = extensions.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);
    Model::MakeDefault();
    LodModel::MakeDefault();
    return true;
}
Renderer::UICtx Renderer::getUICtx() const {
    auto queueFamily = physicalInfo->queueIndices.graphicsFamily.value();
    VkQueue queue;
    vkGetDeviceQueue(device, queueFamily, 0, &queue);
    return {instance, physicalInfo->device, device, queueFamily, queue, (uint32_t)swapChain->getImageViews().size(), pipeline->getUIRenderPass()};
}

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

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

    if (auto newImage = swapChain->acquireNextImage(imageAvailableSemaphores[currentFrame], inFlightFences[currentFrame])) {
        currentImage = newImage.value();
        VoxelUBO ubo{};
        ubo.proj = ProjectionMatrix;
        ubo.view = ViewMatrix;
        ubo.fog = options.clear_color;
        ubo.fog.a = FogDepth;
        ubo.lightInv = LightInvDir;
        commandCenter->startRecording(currentImage, pipeline->getRenderPass(),
            getInfos().swapDetails.capabilities.currentExtent, ubo);
        ShortIndexedVertexBuffer::ClearUnused(currentImage);
    } else {
        recreateSwapChain();
        beginFrame();
    }
}

std::function<size_t(render::LodModel *const, glm::mat4, glm::vec4, float)> Renderer::beginWorldPass(bool solid) {
    assert(currentImage < swapChain->getImageViews().size());

    auto& pass = solid ? pipeline->getWorldPass() : pipeline->getTransparentWorldPass();
    commandCenter->startWorldPass(currentImage, pass);
    return [&](render::LodModel *const rBuffer, glm::mat4 model, glm::vec4 sphere, float curv) {
        auto buffer = dynamic_cast<render::vk::LodModel *const>(rBuffer);
        buffer->setLastUse(currentImage);
        UniqueCurvaturePush push{};
        push.model = model;
        push.sphereProj = sphere;
        push.curvature = curv;
        return commandCenter->recordModel(currentImage, pass, push, buffer);
    };
}

std::function<size_t(render::Model *const, const std::vector<glm::mat4> &)> Renderer::beginEntityPass() {
    assert(currentImage < swapChain->getImageViews().size());

    auto& pass = pipeline->getEntitiesPass();
    commandCenter->startEntityPass(currentImage, pass);
    return [&](render::Model *const rBuffer, const std::vector<glm::mat4>& models) {
        auto buffer = dynamic_cast<render::vk::Model *const>(rBuffer);
        buffer->setLastUse(currentImage);
        return commandCenter->recordModels(currentImage, pass, models, buffer);
    };
    return [](render::Model *const, const std::vector<glm::mat4> &) { return 0; };
}

std::function<size_t(glm::mat4, world::action::Shape, glm::vec4)> Renderer::beginIndicatorPass() {
    assert(currentImage < swapChain->getImageViews().size());

    auto &pass = pipeline->getIndicPass();
    commandCenter->startIndicPass(currentImage, pass);
    return [&](glm::mat4 model, world::action::Shape shape, glm::vec4 color) {
        ModelColorPush push{};
        push.model = model;
        push.color = color;
        return commandCenter->recordIndicator(currentImage, pass, push, shape == world::action::Shape::Cube);
    };
}

void Renderer::postProcess() {
    commandCenter->recordPostprocess(currentImage, pipeline->getSkyPass(), options.skybox);
}
void Renderer::recordUI(std::function<void(VkCommandBuffer)> call) {
    return commandCenter->recordUI(currentImage, pipeline->getUIRenderPass(),
        getInfos().swapDetails.capabilities.currentExtent, call);
}
void Renderer::endFrame() {
    commandCenter->submitGraphics(currentImage, imageAvailableSemaphores[currentFrame],
        renderFinishedSemaphores[currentFrame], inFlightFences[currentFrame]);
}

void Renderer::swapBuffer(Window&) {
    if(!swapChain->presentImage(currentImage, 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 render::passOptions& opt) {
    options.voxel = opt;
    recreateSwapChain();
}
void Renderer::reloadTextures(const std::string& textures, float mipmap, int anisotropy) {
    vkDeviceWaitIdle(device);
    commandCenter->free();
    commandCenter->loadAtlases(textures, anisotropy, mipmap);
    commandCenter->allocate(swapChain->getImageViews(), *pipeline.get(), getInfos().device,
        getInfos().swapDetails.capabilities.currentExtent);
}
void Renderer::setFillMode(bool wireframe) {
    options.wireframe = wireframe;
    recreateSwapChain();
}
void Renderer::setVSync(bool vSync) {
    if (vSync != getInfos().vSync) {
        physicalInfo->vSync = vSync;
        recreateSwapChain();
    }
}

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

void Renderer::setClearColor(glm::vec4 c) {
    options.clear_color = c;
}