#include "vulkan_buffers.h"

#include "src/vulkan_commons.h"
#include "vulkan_errors.h"
#include "locale_vulkan.h"
#include "utilities/logger.h"
#include "utilities/rse_errors_common.h"
#include "rse_math.h"
#include "rse_vulkan_commands.h"

#include <stdint.h>
#include <string.h>

#define MAX_VERTEX_BUFFER_SIZE 33554432 /* 32 MB*/

/* GPU visible buffer */
size_t g_vertex_offset = 0;
struct rse_vulkan_buffer_t g_vertex_buffer;
struct rse_vulkan_buffer_t g_index_buffer;
struct rse_vulkan_buffer_t g_instance_buffers[MAX_MESH_NUMBER];
struct rse_vulkan_buffer_t g_draw_indirect_command_buffer;

/**
 * @brief Copy one buffer's data to another
 *
 * @param src Source buffer
 * @param dst Destination buffer
 * @param size Size of buffer to copy
 */
static void copy_buffer(VkBuffer src, VkBuffer dst, VkDeviceSize size, VkDeviceSize dest_offset)
{
    /* Vulkan buffers can only be copied using command buffers */
    VkBufferCopy copy_region = {};
    VkCommandBuffer command_buffer = {0};

    command_buffer = rse_begin_single_time_command();

    copy_region.srcOffset = 0U;
    copy_region.dstOffset = dest_offset;
    copy_region.size = size;

    vkCmdCopyBuffer(command_buffer, src, dst, 1, &copy_region);

    rse_end_single_time_comands(command_buffer);
}


/**
 * @brief Create a buffer holding all vertex data
 *
 * @return rse_err_t RSE_ERROR_NO_ERROR on success. Possible errors:
 *              VULKAN_ERROR_BUFFER_CREATION_FAILED
 *              VULKAN_ERROR_VERTEX_BUFFER_MAPPING_FAILED
 */
static rse_err_t create_vertex_buffer()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize buffer_size;

    buffer_size = MAX_VERTEX_BUFFER_SIZE;

    /* Create Vertex Buffer*/
    STATUS_CHECK(create_buffer(buffer_size,
                                   VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
                                   VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
                                   0, /* Will not be mapped with vmaMapMemory */
                                   &g_vertex_buffer));

    g_vertex_buffer.allocated_size = 0;
    return RSE_ERROR_NO_ERROR;
}

/**
 * @brief Create an index buffer, connected with vertex buffer
 *
 * @return rse_err_t RSE_ERROR_NO_ERROR on success. Possible errors:
 *              VULKAN_ERROR_BUFFER_CREATION_FAILED
 */
static rse_err_t create_index_buffer()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize buffer_size;

    buffer_size = MAX_VERTEX_BUFFER_SIZE;

    /* Create Index Buffer*/
    STATUS_CHECK(create_buffer(buffer_size, VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
                                   VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
                                   0, /* Will not be mapped with vmaMapMemory */
                                   &g_index_buffer));

    g_index_buffer.allocated_size = 0;
    return RSE_ERROR_NO_ERROR;
}

static rse_err_t create_instance_buffers()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    size_t iter = 0U;
    VkDeviceSize buffer_size;

    buffer_size = sizeof(struct rse_instance_data_t) * MAX_INSTANCE_NUMBER;

    for (iter = 0; iter < MAX_MESH_NUMBER; ++iter) {
        /* Create Instance Buffer */
        STATUS_CHECK(create_buffer(buffer_size,
                                    VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
                                    VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE, 0, /* Will not be mapped with vmaMapMemory */
                                    &g_instance_buffers[iter]));

        g_instance_buffers[iter].allocated_size = 0;

    }

    STATUS_CHECK(create_buffer(sizeof(VkDrawIndexedIndirectCommand) * MAX_MESH_NUMBER,
                                VK_BUFFER_USAGE_TRANSFER_DST_BIT
                                | VK_BUFFER_USAGE_STORAGE_BUFFER_BIT
                                | VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT,
                                VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                                VMA_ALLOCATION_CREATE_MAPPED_BIT | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                                &g_draw_indirect_command_buffer));
    return RSE_ERROR_NO_ERROR;
}

/**
 * @brief Create a Uniform Buffers
 *
 * @return rse_err_t RSE_ERROR_NO_ERROR on success.
 */
static rse_err_t create_uniform_buffer()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize buffer_size = sizeof(struct rse_uniform_buffer_object_t);

    for (size_t i = 0; i < SWAP_BUFFER_COUNT; i++) {
        STATUS_CHECK(create_buffer(buffer_size, VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT, VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                     VMA_ALLOCATION_CREATE_MAPPED_BIT | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                     &g_vulkan_state.uniform_buffers[i]));
    }

    return RSE_ERROR_NO_ERROR;
}

rse_err_t create_buffer(const VkDeviceSize size,
                    VkBufferUsageFlags buffer_usage,
                    VmaMemoryUsage memory_usage,
                    const VmaAllocationCreateFlags allocation_flags,
                    struct rse_vulkan_buffer_t* buffer)
{
    VkBufferCreateInfo vertex_buffer_info = {};
    VmaAllocationCreateInfo create_info = {};

    vertex_buffer_info.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    vertex_buffer_info.pNext = NULL;
    vertex_buffer_info.flags = 0U;
    vertex_buffer_info.size = size;
    vertex_buffer_info.usage = buffer_usage;
    vertex_buffer_info.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    vertex_buffer_info.queueFamilyIndexCount = 0U;
    vertex_buffer_info.pQueueFamilyIndices = NULL;

    create_info.flags = allocation_flags;
    create_info.usage = memory_usage;
    create_info.memoryTypeBits = 0U;
    create_info.requiredFlags = 0U;
    create_info.preferredFlags = 0U;
    create_info.pool = VK_NULL_HANDLE;
    create_info.pUserData = VK_NULL_HANDLE;
    create_info.priority = 0.0f;

    if (VK_SUCCESS != vmaCreateBuffer(g_vulkan_state.allocator,
                                      &vertex_buffer_info,
                                      &create_info,
                                      &buffer->buffer,
                                      &buffer->allocation,
                                      &buffer->allocation_info)) {
        LOGF(vulkan_messages[VULKAN_BUFFER_CREATION_FAILED]);
        return VULKAN_ERROR_BUFFER_CREATION_FAILED;
    }

    return RSE_ERROR_NO_ERROR;
}

void destroy_buffer(struct rse_vulkan_buffer_t* buffer)
{
    vmaDestroyBuffer(g_vulkan_state.allocator, buffer->buffer, buffer->allocation);
}


void update_uniform_buffers()
{
    struct rse_uniform_buffer_object_t ubo = {};
    struct vec3_t eye = {0};
    struct vec3_t center = {0};
    struct vec3_t up = {0};
    struct vec3_t rotation_vec = {1.0f, 0.0f, 0.0f};
    float fovy;
    float z_near;
    float z_far;

    eye.z = 5.0f;

    up.y = 1.0f;

    fovy = rse_math_deg_to_radians(45.0f);
    z_near = 0.1f;
    z_far = 20.0f;

    ubo.model = rse_math_rotate(rse_math_uniform_mat4(), rse_math_deg_to_radians(0.0f), rotation_vec);
    ubo.view = rse_math_look_at(&eye, &center, &up);
    ubo.proj = rse_math_perspective(fovy,
                            g_vulkan_state.swapchain_extent.width / (float)g_vulkan_state.swapchain_extent.height,
                            z_near,
                            z_far);
    memcpy(g_vulkan_state.uniform_buffers[g_vulkan_state.current_frame].allocation_info.pMappedData, &ubo, sizeof(ubo));
}

rse_err_t rse_add_vertices(uint16_t mesh_id,
                         const struct rse_vertex_t* vertices,
                         const uint16_t* indices,
                         struct rse_buffer_data_t* vertex_data,
                         struct rse_buffer_data_t* index_data)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDrawIndexedIndirectCommand draw_indirect_command;
    struct rse_vulkan_buffer_t staging_buffer;
    size_t vertices_size = sizeof(vertices[0]) * vertex_data->count;
    size_t indices_size = sizeof(indices[0]) * index_data->count;

    struct rse_vulkan_buffer_t* index_buffer = &g_index_buffer;

    /* Creating staging buffer*/
    STATUS_CHECK(create_buffer(MAX_VERTEX_BUFFER_SIZE,
                                   VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                                   VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                                   VMA_ALLOCATION_CREATE_MAPPED_BIT
                                   | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                                   &staging_buffer));

    /* Fill staging buffer with vertex data */
    memset(staging_buffer.allocation_info.pMappedData, 0, MAX_VERTEX_BUFFER_SIZE);
    memcpy(staging_buffer.allocation_info.pMappedData, vertices, vertices_size);
    copy_buffer(staging_buffer.buffer,
               g_vertex_buffer.buffer,
               vertices_size,
               g_vertex_buffer.allocated_size);

    vertex_data->buffer_offset = g_vertex_offset;
    g_vertex_offset += vertex_data->count;
    g_vertex_buffer.allocated_size += vertices_size;

    /* Fill staging buffer with index data */
    memset(staging_buffer.allocation_info.pMappedData, 0, indices_size);
    memcpy(staging_buffer.allocation_info.pMappedData, indices, indices_size);
    copy_buffer(staging_buffer.buffer,
               index_buffer->buffer,
               indices_size,
               index_buffer->allocated_size);
    index_data->buffer_offset = index_buffer->allocated_size;
    index_buffer->allocated_size += indices_size;

    vmaDestroyBuffer(g_vulkan_state.allocator, staging_buffer.buffer, staging_buffer.allocation);

    draw_indirect_command.indexCount = 6;
    draw_indirect_command.instanceCount = 0;
    draw_indirect_command.firstIndex = 0;
    draw_indirect_command.vertexOffset = vertex_data->buffer_offset;
    draw_indirect_command.firstInstance = 0;

    memcpy(g_draw_indirect_command_buffer.allocation_info.pMappedData + (mesh_id * sizeof(VkDrawIndexedIndirectCommand)), &draw_indirect_command, sizeof(VkDrawIndexedIndirectCommand));

    return RSE_ERROR_NO_ERROR;
}

rse_err_t update_mesh_instances(uint16_t mesh_id, struct rse_instance_data_t* instance_data)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDrawIndexedIndirectCommand draw_indirect_command;
    struct rse_vulkan_buffer_t staging_buffer;
    struct rse_vulkan_buffer_t* instance_buffer = &g_instance_buffers[mesh_id];
    size_t instance_size = sizeof(struct rse_instance_data_t);

    /* Creating staging buffer*/
    STATUS_CHECK(create_buffer(instance_size,
                                   VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                                   VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                                   VMA_ALLOCATION_CREATE_MAPPED_BIT
                                   | VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT,
                                   &staging_buffer));

    /* Copy transformation information about instance */
    memcpy(staging_buffer.allocation_info.pMappedData, instance_data, instance_size);
    copy_buffer(staging_buffer.buffer,
               instance_buffer->buffer,
               instance_size,
               instance_buffer->allocated_size);

    vmaDestroyBuffer(g_vulkan_state.allocator, staging_buffer.buffer, staging_buffer.allocation);

    instance_buffer->allocated_size += instance_size;

    memcpy(&draw_indirect_command, g_draw_indirect_command_buffer.allocation_info.pMappedData + (mesh_id * sizeof(VkDrawIndexedIndirectCommand)), sizeof(VkDrawIndexedIndirectCommand));
    draw_indirect_command.instanceCount++;
    memcpy(g_draw_indirect_command_buffer.allocation_info.pMappedData + (mesh_id * sizeof(VkDrawIndexedIndirectCommand)), &draw_indirect_command, sizeof(VkDrawIndexedIndirectCommand));

    return RSE_ERROR_NO_ERROR;
}

/**
 * @brief Records commands for provided swapchain framebuffer
 *
 * @param command_buffer Command buffer, that will hold commands
 * @param image_index Image index
 * @return rse_err_t RSE_ERROR_NO_ERROR on success. Possible errors:
 *             VULKAN_ERROR_RECORD_COMMAND_BUFFER_FAILED
 */
rse_err_t record_command_buffer(uint32_t image_index)
{
    VkCommandBuffer command_buffer = g_vulkan_state.command_buffers[g_vulkan_state.current_frame];
    VkDescriptorSet descriptor_sets[] = {g_vulkan_state.descriptor_sets[g_vulkan_state.current_frame],
                                        g_vulkan_state.descriptor_sets_bindless};
    VkCommandBufferBeginInfo begin_info = {};
    VkRenderPassBeginInfo render_pass_info = {};
    VkViewport viewport = {};
    VkRect2D scissor = {};
    VkDeviceSize offsets[] = {0};
    VkClearValue clear_values[2] = {}; //For color and depth stencil
    size_t i = 0U;

    clear_values[0].color.float32[0] = 0.0f;
    clear_values[0].color.float32[1] = 0.0f;
    clear_values[0].color.float32[2] = 0.0f;

    clear_values[1].depthStencil.depth = 1.0f;
    clear_values[1].depthStencil.stencil = 0.0f;

    begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    begin_info.flags = 0;
    begin_info.pInheritanceInfo = NULL;

    if (vkBeginCommandBuffer(command_buffer, &begin_info) != VK_SUCCESS) {
        LOGF(vulkan_messages[VULKAN_RECORD_COMMAND_BEGIN_FAILED]);
        return VULKAN_ERROR_RECORD_COMMAND_BEGIN_FAILED;
    }

    render_pass_info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    render_pass_info.renderPass = g_vulkan_state.render_pass;
    render_pass_info.framebuffer = g_vulkan_state.swapchain_framebuffers[image_index];
    render_pass_info.renderArea.offset.x = 0;
    render_pass_info.renderArea.offset.y = 0;
    render_pass_info.renderArea.extent = g_vulkan_state.swapchain_extent;
    render_pass_info.clearValueCount = 2;
    render_pass_info.pClearValues = clear_values;

    vkCmdBeginRenderPass(command_buffer, &render_pass_info, VK_SUBPASS_CONTENTS_INLINE);
    vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, g_vulkan_state.graphics_pipeline);

    viewport.x = 0.0f;
    viewport.y = 0.0f;
    viewport.width = (float)(g_vulkan_state.swapchain_extent.width);
    viewport.height = (float)(g_vulkan_state.swapchain_extent.height);
    viewport.minDepth = 0.0f;
    viewport.maxDepth = 1.0f;
    vkCmdSetViewport(command_buffer, 0, 1, &viewport);

    scissor.offset.x = 0;
    scissor.offset.y = 0;
    scissor.extent = g_vulkan_state.swapchain_extent;
    vkCmdSetScissor(command_buffer, 0, 1, &scissor);

    vkCmdBindVertexBuffers(command_buffer, 0, 1, &g_vertex_buffer.buffer, offsets);
    vkCmdBindIndexBuffer(command_buffer, g_index_buffer.buffer, 0, VK_INDEX_TYPE_UINT16);
    vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, g_vulkan_state.pipeline_layout, 0, 2,
                            descriptor_sets, 0, NULL);

    for (i = 0; i < 2; ++i) {
        VkBuffer instance_buffer = {g_instance_buffers[i].buffer};

        vkCmdBindVertexBuffers(command_buffer, 1, 1, &instance_buffer, offsets);

        vkCmdDrawIndexedIndirect(command_buffer,
                                g_draw_indirect_command_buffer.buffer,
                                sizeof(VkDrawIndexedIndirectCommand) * i,
                                1,
                                0); /* Can be 0 since we are not doing multiple draws each loop */
    }

    vkCmdEndRenderPass(command_buffer);

    if (VK_SUCCESS != vkEndCommandBuffer(command_buffer)) {
        LOGF(vulkan_messages[VULKAN_RECORD_COMMAND_BUFFER_FAILED]);
        return VULKAN_ERROR_RECORD_COMMAND_BUFFER_FAILED;
    }

    return RSE_ERROR_NO_ERROR;
}

rse_err_t create_buffers()
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    STATUS_CHECK(create_vertex_buffer());
    STATUS_CHECK(create_index_buffer());
    STATUS_CHECK(create_instance_buffers());
    STATUS_CHECK(create_uniform_buffer());

    return status;
}

void reset_command_buffer()
{
    vkResetCommandBuffer(g_vulkan_state.command_buffers[g_vulkan_state.current_frame], /*VkCommandBufferResetFlagBits*/ 0);
}


void destroy_buffers()
{
    size_t i = 0U;

    for (i = 0; i < SWAP_BUFFER_COUNT; ++i) {
        vmaDestroyBuffer(g_vulkan_state.allocator,
                         g_vulkan_state.uniform_buffers[i].buffer,
                         g_vulkan_state.uniform_buffers[i].allocation);
    }

    vmaDestroyBuffer(g_vulkan_state.allocator, g_vertex_buffer.buffer, g_vertex_buffer.allocation);
    vmaDestroyBuffer(g_vulkan_state.allocator, g_index_buffer.buffer, g_index_buffer.allocation);
    vmaDestroyBuffer(g_vulkan_state.allocator, g_draw_indirect_command_buffer.buffer,  g_draw_indirect_command_buffer.allocation);

    for(i = 0; i < MAX_MESH_NUMBER; ++i) {
        vmaDestroyBuffer(g_vulkan_state.allocator, g_instance_buffers[i].buffer, g_instance_buffers[i].allocation);
    }
}