/**
 *
 * @file src/render/vulkan/Vulkan_window_renderer.cpp
 *
 * @copyright 2008-2024 Bordeaux INP, CNRS (LaBRI UMR 5800), Inria,
 *                      Univ. Bordeaux. All rights reserved.
 *
 * @author Camille Ordronneau
 * @author Johnny Jazeix
 * @author Nolan Bredel
 * @author Lucas Guedon
 * @author Augustin Gauchet
 *
 * @date 2024-07-17
 */

#include "common/common.hpp"
#include "common/Info.hpp"
#include "Vulkan_window_renderer.hpp"
#include "common/common.hpp"
#include <QVulkanFunctions>
#include <QFile>
#include <QEventLoop>

/*!
 * \brief Size of the data of the uniform variable
 */
static const int UNIFORM_DATA_SIZE = 16 * 2 * sizeof(float);

/*!
 * \brief The number of uniform buffers
 */
static const int UNIFORM_BUFFER_COUNT = 3;
/*!
 * \brief Finds the smallest multiple of byte_align greater than value
 */
static inline VkDeviceSize aligned(VkDeviceSize value, VkDeviceSize byte_align) {
    byte_align--;
    return (value + byte_align) & ~byte_align;
}

Vulkan_window_renderer::Vulkan_window_renderer(Interface_graphic *interface_graphic, QVulkanWindow *w) :
    _interface_graphic(interface_graphic), _window(w), _dev_funcs(nullptr) {
}

VkShaderModule Vulkan_window_renderer::create_shader(const QString &shader_name) {
    QFile file(shader_name);
    if (!file.open(QIODevice::ReadOnly)) {
        qWarning("Failed to read shader %s", qPrintable(shader_name));
        return VK_NULL_HANDLE;
    }
    QByteArray binary_content = file.readAll();
    file.close();
    // Initialize shader information
    VkShaderModuleCreateInfo shaderInfo = {
        .sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO,
        .codeSize = (size_t)binary_content.size(),
        .pCode = reinterpret_cast<const uint32_t *>(binary_content.constData())
    };
    VkShaderModule shaderModule;
    VkResult err = _dev_funcs->vkCreateShaderModule(_window->device(), &shaderInfo, nullptr, &shaderModule);
    if (err != VK_SUCCESS) {
        qWarning("Failed to create shader module: %d", err);
        return VK_NULL_HANDLE;
    }
    return shaderModule;
}

void Vulkan_window_renderer::create_descriptor_pool() {
    VkDevice dev = _window->device();
    uint32_t concurrent_frame_count = (uint32_t)_window->concurrentFrameCount();
    // Create the descriptor pool and initialize it
    VkDescriptorPoolSize desc_pool_sizes = {
        .type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        .descriptorCount = concurrent_frame_count * UNIFORM_BUFFER_COUNT
    };
    VkDescriptorPoolCreateInfo desc_pool_info = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
        .maxSets = concurrent_frame_count * UNIFORM_BUFFER_COUNT,
        .poolSizeCount = 1,
        .pPoolSizes = &desc_pool_sizes
    };
    VkResult err = _dev_funcs->vkCreateDescriptorPool(dev, &desc_pool_info, nullptr, &_desc_pool);
    if (err != VK_SUCCESS)
        qFatal("Failed to create descriptor pool: %d", err);
    VkDescriptorSetLayoutBinding layout_binding = {
        .binding = 0,
        .descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        .descriptorCount = 1,
        .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
        .pImmutableSamplers = nullptr
    };
    VkDescriptorSetLayoutCreateInfo desc_layout_info = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .bindingCount = 1,
        .pBindings = &layout_binding
    };
    err = _dev_funcs->vkCreateDescriptorSetLayout(dev, &desc_layout_info, nullptr, &_desc_set_layout);
    if (err != VK_SUCCESS)
        qFatal("Failed to create descriptor set layout: %d", err);

    // Initialise unifrom buffers using the pool
    VkDeviceSize min_align = _window->physicalDeviceProperties()->limits.minUniformBufferOffsetAlignment;
    const VkDeviceSize uniform_alloc_size = aligned(UNIFORM_DATA_SIZE, min_align);
    VkDescriptorSetAllocateInfo desc_set_alloc_info = {
        .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
        .pNext = nullptr,
        .descriptorPool = _desc_pool,
        .descriptorSetCount = 1,
        .pSetLayouts = &_desc_set_layout
    };
    _buf_container_uniform.init_uniform(&desc_set_alloc_info, concurrent_frame_count, uniform_alloc_size);
    _buf_state_uniform.init_uniform(&desc_set_alloc_info, concurrent_frame_count, uniform_alloc_size);
    _buf_static_uniform.init_uniform(&desc_set_alloc_info, concurrent_frame_count, uniform_alloc_size);
}

void Vulkan_window_renderer::create_pipeline() {
    VkDevice dev = _window->device();

    // Vertex layout definition
    VkVertexInputBindingDescription vertex_binding_desc = {
        .binding = 0,
        .stride = sizeof(Vertex),
        .inputRate = VK_VERTEX_INPUT_RATE_VERTEX
    };
    VkVertexInputAttributeDescription vertex_attr_desc[] = {
        { // position
          .location = 0,
          .binding = 0,
          .format = VK_FORMAT_R32G32_SFLOAT,
          .offset = offsetof(Vertex, x) },
        { // color
          .location = 1,
          .binding = 0,
          .format = VK_FORMAT_R32G32B32_SFLOAT,
          .offset = offsetof(Vertex, r) }
    };
    VkPipelineVertexInputStateCreateInfo vertex_input_info = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .vertexBindingDescriptionCount = 1,
        .pVertexBindingDescriptions = &vertex_binding_desc,
        .vertexAttributeDescriptionCount = 2,
        .pVertexAttributeDescriptions = vertex_attr_desc
    };

    // Pipeline cache
    VkPipelineCacheCreateInfo pipeline_cache_info = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO
    };
    VkResult err = _dev_funcs->vkCreatePipelineCache(dev, &pipeline_cache_info, nullptr, &_pipeline_cache);
    if (err != VK_SUCCESS)
        qFatal("Failed to create pipeline cache: %d", err);
    // Pipeline layout
    VkPipelineLayoutCreateInfo pipeline_layout_info = {
        .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
        .setLayoutCount = 1,
        .pSetLayouts = &_desc_set_layout
    };
    err = _dev_funcs->vkCreatePipelineLayout(dev, &pipeline_layout_info, nullptr, &_pipeline_layout);
    if (err != VK_SUCCESS)
        qFatal("Failed to create pipeline layout: %d", err);
    // Shaders
    VkShaderModule vert_shader_module = create_shader(QStringLiteral("color_vert.spv"));
    VkShaderModule frag_shader_module = create_shader(QStringLiteral("color_frag.spv"));

    _triangle_pipeline.init(dev, _dev_funcs, _window, vert_shader_module, frag_shader_module, _pipeline_cache, _pipeline_layout, &vertex_input_info, VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST);
    _line_pipeline.init(dev, _dev_funcs, _window, vert_shader_module, frag_shader_module, _pipeline_cache, _pipeline_layout, &vertex_input_info, VK_PRIMITIVE_TOPOLOGY_LINE_LIST);

    if (vert_shader_module)
        _dev_funcs->vkDestroyShaderModule(dev, vert_shader_module, nullptr);
    if (frag_shader_module)
        _dev_funcs->vkDestroyShaderModule(dev, frag_shader_module, nullptr);
}

void Vulkan_window_renderer::initResources() {
    VkDevice dev = _window->device();
    _dev_funcs = _window->vulkanInstance()->deviceFunctions(dev);

    uint32_t memory_index = _window->hostVisibleMemoryIndex();

    buf_container_vertex.init(dev, _dev_funcs, memory_index);
    buf_container_vertex.alloc_buffer((VkDeviceSize)0);
    buf_state_vertex.init(dev, _dev_funcs, memory_index);
    buf_state_vertex.alloc_buffer((VkDeviceSize)0);
    buf_ruler.init(dev, _dev_funcs, memory_index);
    buf_ruler.alloc_buffer((VkDeviceSize)0);
    buf_counter.init(dev, _dev_funcs, memory_index);
    buf_counter.alloc_buffer((VkDeviceSize)0);
    _buf_container_uniform.init(dev, _dev_funcs, memory_index);
    _buf_state_uniform.init(dev, _dev_funcs, memory_index);
    _buf_static_uniform.init(dev, _dev_funcs, memory_index);

    create_descriptor_pool();

    create_pipeline();

    // Initialise view matrices
    container_model_view.setToIdentity();
    state_model_view.setToIdentity();

    _interface_graphic->get_console()->waitGUIInit->quit();
}

void Vulkan_window_renderer::initSwapChainResources() {
    QSize frame_size = _window->swapChainImageSize();
    Info::Screen::width = 100;
    Info::Screen::height = 100;
    // Projection matrix
    _proj = _window->clipCorrectionMatrix(); // adjust for Vulkan-OpenGL clip space differences
    _proj.ortho(0.0f, 100.0f, 100.0f, 0.0f, -100.0f, 100.0f);
}

void Vulkan_window_renderer::releaseSwapChainResources() { }

void Vulkan_window_renderer::releaseResources() {
    VkDevice dev = _window->device();
    _triangle_pipeline.free_pipeline();
    _line_pipeline.free_pipeline();

    if (_pipeline_layout) {
        _dev_funcs->vkDestroyPipelineLayout(dev, _pipeline_layout, nullptr);
        _pipeline_layout = VK_NULL_HANDLE;
    }

    if (_pipeline_cache) {
        _dev_funcs->vkDestroyPipelineCache(dev, _pipeline_cache, nullptr);
        _pipeline_cache = VK_NULL_HANDLE;
    }

    if (_desc_set_layout) {
        _dev_funcs->vkDestroyDescriptorSetLayout(dev, _desc_set_layout, nullptr);
        _desc_set_layout = VK_NULL_HANDLE;
    }

    if (_desc_pool) {
        _dev_funcs->vkDestroyDescriptorPool(dev, _desc_pool, nullptr);
        _desc_pool = VK_NULL_HANDLE;
    }

    buf_container_vertex.free_buffer();
    buf_state_vertex.free_buffer();
    buf_ruler.free_buffer();
    buf_counter.free_buffer();

    _buf_container_uniform.free_buffer();
    _buf_state_uniform.free_buffer();
    _buf_static_uniform.free_buffer();
}

void Vulkan_window_renderer::draw_vertices(Vk_uniform_buffer *uniform_buf, Vk_vertex_buffer *vertex_buf) {
    VkDevice dev = _window->device();
    VkCommandBuffer vk_command_buffer = _window->currentCommandBuffer();
    uniform_buf->bind_uniform_buffer(vk_command_buffer, _window->currentFrame(), _pipeline_layout);
    vertex_buf->bind_vertex_buffer(&vk_command_buffer);
    vertex_buf->draw(&vk_command_buffer);
}

void Vulkan_window_renderer::startNextFrame() {
    // Send uniform values
    _buf_container_uniform.set_variable(_window->currentFrame(), 16 * sizeof(float), 16 * sizeof(float), _proj.constData());
    _buf_container_uniform.set_variable(_window->currentFrame(), 0, 16 * sizeof(float), container_model_view.constData());
    _buf_state_uniform.set_variable(_window->currentFrame(), 16 * sizeof(float), 16 * sizeof(float), _proj.constData());
    _buf_state_uniform.set_variable(_window->currentFrame(), 0, 16 * sizeof(float), state_model_view.constData());
    _buf_static_uniform.set_variable(_window->currentFrame(), 16 * sizeof(float), 16 * sizeof(float), _proj.constData());

    VkCommandBuffer vk_command_buffer = _window->currentCommandBuffer();
    QSize frame_size = _window->swapChainImageSize();
    // Start render pass
    VkClearColorValue clear_color = { .float32 = { 0.501f, 0.501f, 0.549f, 1.0f } };
    VkClearDepthStencilValue clearDS = { .depth = 1.0f, .stencil = 0 };
    VkClearValue clear_values[] = {
        { .color = clear_color },
        { .depthStencil = clearDS },
        { .color = clear_color }
    };
    VkRenderPassBeginInfo rp_begin_info = {
        .sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO,
        .renderPass = _window->defaultRenderPass(),
        .framebuffer = _window->currentFramebuffer(),
        .renderArea = { .extent = { .width = (uint32_t)frame_size.width(), .height = (uint32_t)frame_size.height() } },
        .clearValueCount = _window->sampleCountFlagBits() > VK_SAMPLE_COUNT_1_BIT ? (uint32_t)3 : (uint32_t)2,
        .pClearValues = clear_values
    };
    _dev_funcs->vkCmdBeginRenderPass(vk_command_buffer, &rp_begin_info, VK_SUBPASS_CONTENTS_INLINE);

    // Setup render area size
    VkViewport viewport = {
        .x = 0.0f,
        .y = 0.0f,
        .width = (float)frame_size.width(),
        .height = (float)frame_size.height(),
        .minDepth = 0.0f,
        .maxDepth = 1.0f
    };
    _dev_funcs->vkCmdSetViewport(vk_command_buffer, 0, 1, &viewport);
    VkRect2D scissor = {
        .offset = { .x = 0, .y = 0 },
        .extent = { .width = (uint32_t)viewport.width, .height = (uint32_t)viewport.height }
    };
    _dev_funcs->vkCmdSetScissor(vk_command_buffer, 0, 1, &scissor);

    // Prepare triangle rendering
    _triangle_pipeline.bind_pipeline(&vk_command_buffer);
    // Draw containers
    draw_vertices(&_buf_container_uniform, &buf_container_vertex);
    // Draw states
    draw_vertices(&_buf_state_uniform, &buf_state_vertex);

    // Prepare line rendering
    _line_pipeline.bind_pipeline(&vk_command_buffer);
    // Draw counter
    draw_vertices(&_buf_state_uniform, &buf_counter);
    // Draw ruler
    draw_vertices(&_buf_static_uniform, &buf_ruler);

    // End render pass
    _dev_funcs->vkCmdEndRenderPass(vk_command_buffer);

    _window->frameReady();
}
