/**
 * @file descriptor_builder.c
 * @author Piotr Krygier (everyonecancode@gmail.com)
 * @brief
 * @version 0.1
 * @date 2025-03-14
 *
 * @copyright Copyright (c) 2025
 *
 */

#include "descriptor_builder.h"

#include <SDL3/SDL_log.h>
#include <stddef.h>
#include <stdint.h>

#include "src/graphics_context.h"
#include "src/vulkan_commons.h"
#include "utilities/commons.h"
#include "utilities/errors_common.h"
#include "vulkan/vulkan_core.h"

#define POOL_MAX_SETS (32U)

rse_err_t descriptor_pool_add_type(struct graphics_context_t* context, VkDescriptorType type, uint32_t count)
{
    context->descriptor_data.pool_sizes[context->descriptor_data.pool_sizes_count].type = type;
    context->descriptor_data.pool_sizes[context->descriptor_data.pool_sizes_count].descriptorCount = count;

    context->descriptor_data.pool_sizes_count++;
    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_pool_initialize(struct graphics_context_t* context)
{
    VkDescriptorPoolCreateInfo create_info = {0};

    create_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
    create_info.pNext = NULL;
    create_info.flags = 0U;
    create_info.maxSets = POOL_MAX_SETS;
    create_info.poolSizeCount = context->descriptor_data.pool_sizes_count;
    create_info.pPoolSizes = context->descriptor_data.pool_sizes;

    if (vkCreateDescriptorPool(context->vulkan_handles.device,
                               &create_info,
                               NULL,
                               &context->descriptor_data.descriptor_pool) != VK_SUCCESS) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create descriptor pool");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_create_new_set(struct graphics_context_t* context, rse_id_t* set_id)
{
    if (context->descriptor_data.descriptor_sets_count >= DESCRIPTOR_MAX_SETS) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Reached maximum number of available descriptor sets");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    *set_id = context->descriptor_data.descriptor_sets_count++;
    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_add_layout(struct descriptor_set_layout_bindings_t* layout_bindings,
                                uint32_t binding,
                                VkDescriptorType type,
                                uint32_t descriptor_count,
                                VkShaderStageFlags shader_stages)
{
    layout_bindings->layout_bindings[layout_bindings->layout_bindings_count].binding = binding;
    layout_bindings->layout_bindings[layout_bindings->layout_bindings_count].descriptorType = type;
    layout_bindings->layout_bindings[layout_bindings->layout_bindings_count].descriptorCount = descriptor_count;
    layout_bindings->layout_bindings[layout_bindings->layout_bindings_count].stageFlags = shader_stages;

    layout_bindings->layout_bindings_count++;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_set_finish(struct graphics_context_t* context,
                                struct descriptor_set_layout_bindings_t* layout_bindings,
                                rse_id_t set_id)
{
    VkDescriptorSetLayoutCreateInfo create_info = {0};

    create_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    create_info.bindingCount = layout_bindings->layout_bindings_count;
    create_info.pBindings = layout_bindings->layout_bindings;

    if (VK_SUCCESS != vkCreateDescriptorSetLayout(context->vulkan_handles.device,
                                                  &create_info,
                                                  NULL,
                                                  &context->descriptor_data.set_layouts[set_id])) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create descriptor sets layout");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    layout_bindings->layout_bindings_count = 0U;
    rse_memset(layout_bindings->layout_bindings, 0, sizeof(VkDescriptorSetLayoutBinding) * DESCRIPTOR_MAX_SETS);

    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_build_sets(struct graphics_context_t* context)
{
    VkDescriptorSetAllocateInfo alloc_info = {0};

    if (NULL == context) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Provided NULL context");

        return RSE_ERROR_NULL_POINTER;
    }

    if (context->descriptor_data.descriptor_pool == VK_NULL_HANDLE) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU,
                        "Requested descriptor build, but descriptor pool has not been initialized");
    }

    /* Allocate descriptors from the pool */
    alloc_info.descriptorPool = context->descriptor_data.descriptor_pool;
    alloc_info.descriptorSetCount = context->descriptor_data.descriptor_sets_count;
    alloc_info.pSetLayouts = context->descriptor_data.set_layouts;
    alloc_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;

    if (VK_SUCCESS != vkAllocateDescriptorSets(context->vulkan_handles.device,
                                               &alloc_info,
                                               context->descriptor_data.descriptor_sets)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create descriptor sets");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_attach_images(struct graphics_context_t* context,
                                   rse_id_t descriptor_set_id,
                                   rse_id_t* texture_ids,
                                   size_t images_count,
                                   VkSampler sampler,
                                   uint32_t binding)
{
    size_t iter = 0U;
    VkDescriptorImageInfo* image_infos = NULL;
    VkWriteDescriptorSet write_set = {0};

    rse_malloc(image_infos, sizeof(VkDescriptorImageInfo) * images_count);
    for (iter = 0U; iter < images_count; ++iter) {
        image_infos[iter].sampler = sampler;
        image_infos[iter].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        image_infos[iter].imageView = context->texture_images[texture_ids[iter]].image_view;
    }

    write_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_set.pImageInfo = image_infos;
    write_set.dstSet = context->descriptor_data.descriptor_sets[descriptor_set_id];
    write_set.dstBinding = binding;
    write_set.dstArrayElement = 0U;
    write_set.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;  // TODO: Maybe not only samplers?
    write_set.descriptorCount = images_count;

    vkUpdateDescriptorSets(context->vulkan_handles.device, 1U, &write_set, 0, NULL);
    rse_free(image_infos);

    return RSE_ERROR_NO_ERROR;
}

rse_err_t descriptor_attach_buffer(struct graphics_context_t* context,
                                   rse_id_t descriptor_set_id,
                                   struct vulkan_buffer_t* buffer,
                                   uint32_t binding,
                                   VkDescriptorType type)
{
    VkDescriptorBufferInfo buffer_info = {0};
    VkWriteDescriptorSet write_set = {0};

    buffer_info.buffer = buffer->buffer;
    buffer_info.offset = 0U;
    buffer_info.range = buffer->allocated_size;

    write_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_set.pBufferInfo = &buffer_info;
    write_set.dstSet = context->descriptor_data.descriptor_sets[descriptor_set_id];
    write_set.dstBinding = binding;
    write_set.dstArrayElement = 0U;
    write_set.descriptorType = type;
    write_set.descriptorCount = 1;

    vkUpdateDescriptorSets(context->vulkan_handles.device, 1U, &write_set, 0, NULL);

    return RSE_ERROR_NO_ERROR;
}

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

    for (iter = 0U; iter < context->descriptor_data.descriptor_sets_count; ++iter) {
        vkDestroyDescriptorSetLayout(context->vulkan_handles.device, context->descriptor_data.set_layouts[iter], NULL);
    }
    vkDestroyDescriptorPool(context->vulkan_handles.device, context->descriptor_data.descriptor_pool, NULL);
}