#include "vulkan_buffers.h"

#include <cglm/affine-pre.h>
#include <cglm/cam.h>
#include <cglm/cglm.h>
#include <cglm/mat4.h>
#include <cglm/util.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>
#include <vulkan/vulkan_core.h>

#include "src/graphics_context.h"
#include "src/vulkan_commons.h"
#include "utilities/commons.h"
#include "utilities/errors_common.h"
#include "utilities/vector.h"
#include "vulkan_commands.h"

#define MAX_VERTEX_BUFFER_SIZE 33554432 /* 32 MB*/
#define CMD_BUFFER_SIZE (MAX_INSTANCE_NUMBER * sizeof(VkDrawIndexedIndirectCommand))

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

    STATUS_CHECK(begin_single_time_command(context, &command_buffer));

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

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

    STATUS_CHECK(end_single_time_comands(context, command_buffer));

    return status;
}

/**
 * @brief Create a buffer holding all vertex data
 *
 * @param context Graphics context handle
 * @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(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize buffer_size;

    buffer_size = MAX_VERTEX_BUFFER_SIZE;

    /* Create Vertex Buffer*/
    STATUS_CHECK(create_buffer(context,
                               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 */
                               &context->render_targets[RENDER_TARGET_3D].vertex_buffer));

    context->render_targets[RENDER_TARGET_3D].vertex_buffer.allocated_size = 0;

    return status;
}

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

    buffer_size = MAX_VERTEX_BUFFER_SIZE;

    /* Create Index Buffer*/
    STATUS_CHECK(create_buffer(context,
                               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 */
                               &context->render_targets[RENDER_TARGET_3D].index_buffer));

    context->render_targets[RENDER_TARGET_3D].index_buffer.allocated_size = 0;
    return RSE_ERROR_NO_ERROR;
}

/**
 * @brief Create instance buffer alongside draw indirect command buffer
 *
 * @param[in/out] context Graphics context handle
 * @return RSE_ERROR_NO_ERROR on success.
 */
static rse_err_t create_instance_buffers(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkDeviceSize buffer_size;

    buffer_size = sizeof(struct instance_data_t) * MAX_INSTANCE_NUMBER;

    STATUS_CHECK(
        create_buffer(context,
                      buffer_size,
                      VK_BUFFER_USAGE_STORAGE_BUFFER_BIT,
                      VMA_MEMORY_USAGE_AUTO_PREFER_HOST,
                      VMA_ALLOCATION_CREATE_HOST_ACCESS_SEQUENTIAL_WRITE_BIT | VMA_ALLOCATION_CREATE_MAPPED_BIT,
                      &context->render_targets[RENDER_TARGET_3D].instance_buffer));
    /* TODO: Move creation to somewhere else? */
    STATUS_CHECK(create_buffer(context,
                               CMD_BUFFER_SIZE,
                               VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDIRECT_BUFFER_BIT,
                               VMA_MEMORY_USAGE_AUTO_PREFER_DEVICE,
                               0,
                               &context->render_targets[RENDER_TARGET_3D].draw_indirect_command_buffer));

    context->render_targets[RENDER_TARGET_3D].instance_buffer.allocated_size = 0;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t buffer_create_uniform(struct graphics_context_t* context,
                                struct vulkan_buffer_t* buffer,
                                VkDeviceSize buffer_size)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    STATUS_CHECK(
        create_buffer(context,
                      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,
                      buffer));

    context->uniform_buffers.vulkan_buffers[context->uniform_buffers.vulkan_buffers_count++] = *buffer;

    return RSE_ERROR_NO_ERROR;
}

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

    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(context->vulkan_handles.allocator,
                                      &vertex_buffer_info,
                                      &create_info,
                                      &buffer->buffer,
                                      &buffer->allocation,
                                      &buffer->allocation_info)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create a vertex buffer");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    buffer->allocated_size = size;

    return RSE_ERROR_NO_ERROR;
}

void destroy_buffer(struct graphics_context_t* context, struct vulkan_buffer_t* buffer)
{
    vmaDestroyBuffer(context->vulkan_handles.allocator, buffer->buffer, buffer->allocation);
}

rse_err_t add_vertices(struct graphics_context_t* context,
                       uint16_t mesh_id,
                       size_t vertices_count,
                       const struct vertex_t* vertices,
                       size_t indices_count,
                       uint16_t* indices)
{
    // static uint16_t last_index = 0U;
    // uint16_t max_index = 0U;
    // size_t idx = 0U;
    rse_err_t status = RSE_ERROR_NO_ERROR;
    struct vulkan_buffer_t staging_buffer;
    size_t vertices_size = sizeof(vertices[0]) * vertices_count;
    size_t indices_size = sizeof(indices[0]) * indices_count;
    size_t staging_buffer_size = vertices_size > indices_size ? vertices_size : indices_size;

    struct vulkan_buffer_t* vertex_buffer = &context->render_targets[RENDER_TARGET_3D].vertex_buffer;
    struct vulkan_buffer_t* index_buffer = &context->render_targets[RENDER_TARGET_3D].index_buffer;

    /* Creating staging buffer*/
    STATUS_CHECK(
        create_buffer(context,
                      staging_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 */
    memcpy(staging_buffer.allocation_info.pMappedData, vertices, vertices_size);
    copy_buffer(context, staging_buffer.buffer, vertex_buffer->buffer, vertices_size, vertex_buffer->allocated_size);
    vertex_buffer->allocated_size += vertices_size;


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

    destroy_buffer(context, &staging_buffer);

    context->mesh_data.mesh_bucket[mesh_id].is_dirty = true;
    context->mesh_data.mesh_bucket[mesh_id].vertex_count = vertices_count;
    context->mesh_data.mesh_bucket[mesh_id].entities_count = 0;
    context->mesh_data.mesh_bucket[mesh_id].instances_count = 0;
    context->mesh_data.mesh_bucket[mesh_id].index_count = indices_count;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t mesh_add_instance(struct graphics_context_t* context, uint32_t mesh_id, struct instance_data_t* instance_data)
{
    /* For now only store information about instances. Upload them to instance buffer later */
    RSE_VECTOR_PUSH_BACK(context->mesh_data.mesh_bucket[mesh_id].instances_data, *instance_data);
    context->mesh_data.mesh_bucket[mesh_id].instances_count++;
    context->mesh_data.mesh_bucket[mesh_id].is_dirty = true;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t buffers_update(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;
    bool dirty_found = false;
    char* instance_buffer_mapped = NULL;
    size_t iter = 0U;
    size_t buffer_offset = 0U;
    VkDrawIndexedIndirectCommand* draw_indirect_commands_buffer = NULL;
    VkDrawIndexedIndirectCommand* draw_indirect_command = NULL;
    VkDeviceSize draw_indirect_commands_buffer_size =
        sizeof(VkDrawIndexedIndirectCommand) * context->mesh_data.mesh_count;
    struct vulkan_buffer_t staging_buffer;
    struct mesh_t* mesh = NULL;
    uint32_t indices_count = 0U;

    rse_malloc(draw_indirect_commands_buffer, draw_indirect_commands_buffer_size);
    for (iter = 0U; iter < context->mesh_data.mesh_count; ++iter) {
        mesh = &context->mesh_data.mesh_bucket[iter];
        draw_indirect_command = &draw_indirect_commands_buffer[iter];
        /* TODO: I wonder if we save some time here actually by checking for dirty...*/
        if ((dirty_found == true) || (context->mesh_data.mesh_bucket[iter].is_dirty == true)) {
            memcpy((char*)context->render_targets[RENDER_TARGET_3D].instance_buffer.allocation_info.pMappedData +
                       buffer_offset,
                   mesh->instances_data.data,
                   mesh->instances_count * sizeof(struct instance_data_t));
            dirty_found = true;
        }
        mesh->is_dirty = false;
        buffer_offset += mesh->instances_count * sizeof(struct instance_data_t);

        if (iter == 0) {
            draw_indirect_command->firstIndex = 0;
            draw_indirect_command->firstInstance = 0;
            draw_indirect_command->vertexOffset = 0;
        } else {
            draw_indirect_command->firstIndex = indices_count;
            draw_indirect_command->firstInstance = draw_indirect_commands_buffer[iter - 1].firstInstance +
                                                   draw_indirect_commands_buffer[iter - 1].instanceCount;
            draw_indirect_command->vertexOffset = (context->mesh_data.mesh_bucket[iter - 1].vertex_count + draw_indirect_commands_buffer[iter-1].vertexOffset);
        }
        draw_indirect_command->indexCount = mesh->index_count;
        draw_indirect_command->instanceCount = mesh->instances_count;
        indices_count += mesh->index_count;
    }

    /* Copy Draw indirect commands into buffer */
    if (dirty_found == true) {
        STATUS_CHECK(
            create_buffer(context,
                          draw_indirect_commands_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));

        /* Copy command buffer information about instances */
        memcpy(staging_buffer.allocation_info.pMappedData,
               draw_indirect_commands_buffer,
               draw_indirect_commands_buffer_size);
        copy_buffer(context,
                    staging_buffer.buffer,
                    context->render_targets[RENDER_TARGET_3D].draw_indirect_command_buffer.buffer,
                    draw_indirect_commands_buffer_size,
                    0);

        destroy_buffer(context, &staging_buffer);
    }

    /* Update instance data uniform buffer */
    context->render_targets[RENDER_TARGET_3D].instance_buffer.allocated_size = buffer_offset;
    buffer_offset = 0U;
    instance_buffer_mapped =
        (char*)context->render_targets[RENDER_TARGET_3D].instance_buffer.allocation_info.pMappedData;
    for (iter = 0U; iter < context->mesh_data.mesh_count; ++iter) {
        mesh = &context->mesh_data.mesh_bucket[iter];
        rse_memcpy(instance_buffer_mapped + buffer_offset,
                   mesh->instances_data.data,
                   mesh->instances_count * sizeof(struct instance_data_t));
        buffer_offset += mesh->instances_count * sizeof(struct instance_data_t);
    }

    rse_free(draw_indirect_commands_buffer);

    return status;
}

/**
 * @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(struct graphics_context_t* context, uint32_t image_index)
{
    size_t iter = 0U;
    VkCommandBuffer command_buffer = context->vulkan_handles.command_buffers[context->current_frame];
    VkCommandBufferBeginInfo begin_info = {0};
    VkRenderPassBeginInfo render_pass_info = {0};
    VkViewport viewport = {0};
    VkRect2D scissor = {0};
    VkClearValue clear_values[2] = {0}; /* For color and depth stencil */

    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) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to start recording command buffer");
        return RSE_ERROR_INTERNAL_ERROR;
    }

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

    vkCmdBeginRenderPass(command_buffer, &render_pass_info, VK_SUBPASS_CONTENTS_INLINE);
    viewport.x = 0.0f;
    viewport.y = 0.0f;
    viewport.width = (float)(context->swapchain_data.swapchain_extent.width);
    viewport.height = (float)(context->swapchain_data.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 = context->swapchain_data.swapchain_extent;
    vkCmdSetScissor(command_buffer, 0, 1, &scissor);

    for (iter = 0U; iter < context->pipelines_data.pipelines_count; ++iter) {
        context->renderer_data.renderers[iter].render_function(context, iter);
    }

    vkCmdEndRenderPass(command_buffer);

    if (VK_SUCCESS != vkEndCommandBuffer(command_buffer)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to record command buffer. Validate your commands.");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    return RSE_ERROR_NO_ERROR;
}

rse_err_t create_buffers(struct graphics_context_t* context)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    STATUS_CHECK(create_vertex_buffer(context));
    STATUS_CHECK(create_index_buffer(context));
    STATUS_CHECK(create_instance_buffers(context));

    return status;
}

void reset_command_buffer(struct graphics_context_t* context)
{
    vkResetCommandBuffer(context->vulkan_handles.command_buffers[context->current_frame],
                         /*VkCommandBufferResetFlagBits*/ 0);
}

void destroy_buffers(struct graphics_context_t* context)
{
    size_t iter = 0U;

    destroy_buffer(context, &context->render_targets[RENDER_TARGET_3D].vertex_buffer);
    destroy_buffer(context, &context->render_targets[RENDER_TARGET_3D].index_buffer);
    destroy_buffer(context, &context->render_targets[RENDER_TARGET_3D].draw_indirect_command_buffer);
    destroy_buffer(context, &context->render_targets[RENDER_TARGET_3D].instance_buffer);

    for (iter = 0U; iter < context->uniform_buffers.vulkan_buffers_count; ++iter) {
        vmaDestroyBuffer(context->vulkan_handles.allocator,
                         context->uniform_buffers.vulkan_buffers[iter].buffer,
                         context->uniform_buffers.vulkan_buffers[iter].allocation);
    }

    context->uniform_buffers.vulkan_buffers_count = 0;
    for (iter = 0; iter < context->mesh_data.mesh_count; ++iter) {
        rse_free(context->mesh_data.mesh_bucket[iter].instances_data.data);
    }
}