/**
 * @file pipeline_builder.c
 * @author Piotr Krygier (everyonecancode@gmail.com)
 * @brief
 * @version 0.1
 * @date 2024-03-01
 *
 * @copyright Copyright (c) 2024
 *
 */

#include "pipeline_builder.h"

#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#include "SDL3/SDL_log.h"
#include "src/graphics_context.h"
#include "utilities/commons.h"
#include "utilities/errors_common.h"
#include "utilities/file_utils.h"
#include "vulkan/vulkan_core.h"

#define MAX_BINDLESS_RESOURCES 1000

rse_err_t shader_create_module(struct graphics_context_t* context,
                               rse_id_t* shader_module_id,
                               const char* shader_path,
                               const VkShaderStageFlagBits shader_stage)
{
    char* buffer;
    size_t buffer_size;
    VkShaderModuleCreateInfo create_info = {0};

    if (shader_module_id == NULL) {
        SDL_LogError(SDL_LOG_CATEGORY_GPU, "Please provide valid pointer to shader module ID\n");
    }

    /* Check if shader module is valid */
    if (shader_stage >= VK_SHADER_STAGE_FLAG_BITS_MAX_ENUM) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Invalid shader stage in pipeline");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    /* Read shader file */
    file_read(shader_path, &buffer_size, NULL);
    rse_malloc(buffer, buffer_size);
    file_read(shader_path, &buffer_size, buffer);

    create_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    create_info.pNext = NULL;
    create_info.flags = 0U;
    create_info.codeSize = buffer_size;
    create_info.pCode = (const uint32_t*)(buffer);

    if (context->shader_modules.shader_modules_count > MAX_SHADER_MODULES_COUNT) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Reached maximum number of available shader modules");

        return RSE_ERROR_INTERNAL_ERROR;
    }

    if (VK_SUCCESS !=
        vkCreateShaderModule(context->vulkan_handles.device,
                             &create_info,
                             NULL,
                             &context->shader_modules.shader_modules[context->shader_modules.shader_modules_count])) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create a shader");

        return RSE_ERROR_INTERNAL_ERROR;
    }

    context->shader_modules.shaders_stages[context->shader_modules.shader_modules_count] = shader_stage;

    *shader_module_id = context->shader_modules.shader_modules_count++;
    rse_free(buffer);

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_input_assembly(struct pipeline_t* pipeline)
{
    /* For now input assebly is hardcoded */
    pipeline->input_assembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    pipeline->input_assembly.pNext = NULL;
    pipeline->input_assembly.flags = 0U;
    pipeline->input_assembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    pipeline->input_assembly.primitiveRestartEnable = VK_FALSE;

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_viewport_state(struct pipeline_t* pipeline)
{
    /* For now viewport is hardcoded */
    pipeline->viewport_state.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    pipeline->viewport_state.pNext = NULL;
    pipeline->viewport_state.flags = 0U;
    pipeline->viewport_state.viewportCount = 1U;
    pipeline->viewport_state.pViewports = NULL; /* This is set during rendering */
    pipeline->viewport_state.scissorCount = 1U;
    pipeline->viewport_state.pScissors = NULL; /* This is set during rendering */

    if (pipeline->viewport_state.pViewports == NULL || pipeline->viewport_state.viewportCount == 0U) {
        pipeline->required_dynamic_states.dynamic_states[pipeline->required_dynamic_states.dynamic_states_count++] =
            VK_DYNAMIC_STATE_VIEWPORT;
    }

    if (pipeline->viewport_state.pScissors == NULL || pipeline->viewport_state.scissorCount == 0U) {
        pipeline->required_dynamic_states.dynamic_states[pipeline->required_dynamic_states.dynamic_states_count++] =
            VK_DYNAMIC_STATE_SCISSOR;
    }

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_rasterizer(struct pipeline_t* pipeline)
{
    /* For now rasterizer is hardcoded */
    pipeline->rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    pipeline->rasterizer.pNext = NULL;
    pipeline->rasterizer.flags = 0U;
    pipeline->rasterizer.depthClampEnable = VK_FALSE;        /* VK_TRUE to play with shadow maps */
    pipeline->rasterizer.rasterizerDiscardEnable = VK_FALSE; /* Disable/Enable output to framebuffer*/
    pipeline->rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
    pipeline->rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
    pipeline->rasterizer.frontFace = VK_FRONT_FACE_CLOCKWISE;
    pipeline->rasterizer.depthBiasEnable = VK_FALSE;
    pipeline->rasterizer.depthBiasConstantFactor = 0.0f;
    pipeline->rasterizer.depthBiasClamp = 0.0f;
    pipeline->rasterizer.depthBiasSlopeFactor = 0.0f;
    pipeline->rasterizer.lineWidth = 1.0f;

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_multisampling(struct pipeline_t* pipeline)
{
    /* For now multisampling is hardcoded */
    pipeline->multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    pipeline->multisampling.pNext = NULL;
    pipeline->multisampling.flags = 0U;
    pipeline->multisampling.rasterizationSamples = VK_SAMPLE_COUNT_8_BIT;
    pipeline->multisampling.sampleShadingEnable = VK_FALSE;
    pipeline->multisampling.minSampleShading = 1.0f;
    pipeline->multisampling.pSampleMask = NULL;
    pipeline->multisampling.alphaToCoverageEnable = VK_FALSE;
    pipeline->multisampling.alphaToOneEnable = VK_FALSE;

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_color_blending(struct pipeline_t* pipeline)
{
    /* For now color blending is hardcoded */
    pipeline->color_blend_attachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
    pipeline->color_blend_attachment.blendEnable = VK_TRUE;
    pipeline->color_blend_attachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
    pipeline->color_blend_attachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
    pipeline->color_blend_attachment.colorBlendOp = VK_BLEND_OP_ADD;
    pipeline->color_blend_attachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    pipeline->color_blend_attachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    pipeline->color_blend_attachment.alphaBlendOp = VK_BLEND_OP_ADD;

    pipeline->color_blending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    pipeline->color_blending.pNext = NULL;
    pipeline->color_blending.flags = 0U;
    pipeline->color_blending.logicOpEnable = VK_FALSE;
    pipeline->color_blending.logicOp = VK_LOGIC_OP_COPY;
    pipeline->color_blending.attachmentCount = 1U;
    pipeline->color_blending.pAttachments = &pipeline->color_blend_attachment;
    pipeline->color_blending.blendConstants[0] = 0.0f;
    pipeline->color_blending.blendConstants[1] = 0.0f;
    pipeline->color_blending.blendConstants[2] = 0.0f;
    pipeline->color_blending.blendConstants[3] = 0.0f;

    return RSE_ERROR_NO_ERROR;
}

static rse_err_t pipeline_add_depth_stencil(struct pipeline_t* pipeline)
{
    /* For now depth stencil is hardcoded */
    pipeline->depth_stencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    pipeline->depth_stencil.pNext = NULL;
    pipeline->depth_stencil.flags = 0U;
    pipeline->depth_stencil.depthTestEnable = VK_TRUE;
    pipeline->depth_stencil.depthWriteEnable = VK_TRUE;
    pipeline->depth_stencil.depthCompareOp = VK_COMPARE_OP_LESS;
    pipeline->depth_stencil.depthBoundsTestEnable = VK_FALSE;
    pipeline->depth_stencil.stencilTestEnable = VK_FALSE;
    pipeline->depth_stencil.minDepthBounds = 0.0f;
    pipeline->depth_stencil.maxDepthBounds = 1.0f;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t pipeline_add_descriptor_sets(struct graphics_context_t* context,
                                       struct pipeline_t* pipeline,
                                       rse_id_t descriptor_set_id)
{
    VkPipelineLayoutCreateInfo create_info = {0};
    create_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    create_info.pNext = NULL;
    create_info.flags = 0U;
    create_info.setLayoutCount = 1;  // FIXME: Make sure we can enable more than one layout
    create_info.pSetLayouts = &context->descriptor_data.set_layouts[descriptor_set_id];
    create_info.pushConstantRangeCount = pipeline->push_constant_ranges.push_constant_ranges_count;
    create_info.pPushConstantRanges = pipeline->push_constant_ranges.push_constant_ranges;

    if (VK_SUCCESS !=
        vkCreatePipelineLayout(context->vulkan_handles.device, &create_info, NULL, &pipeline->pipeline_layout)) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create pipeline layout");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    context->pipelines_data.descriptor_set_ids[context->pipelines_data.pipelines_count] = descriptor_set_id;
    context->pipelines_data.pipeline_layouts[context->pipelines_data.pipeline_layouts_count++] =
        pipeline->pipeline_layout;

    return RSE_ERROR_NO_ERROR;
}

/**
 * @brief Clean up pipeline builder. This function should be called after all
 * pipelines are built and no more pipelines are needed. It frees all resources
 * used by pipeline builder.
 *
 * @param[in/out] pipelines List of pipelines to clean up
 */
static void pipeline_infos_cleanup(struct pipeline_infos_t* pipelines)
{
    size_t pipeline_count = pipelines->pipeline_create_infos_count;

    for (size_t i = 0; i < pipeline_count; i++) {
        VkGraphicsPipelineCreateInfo* pipeline_info = &pipelines->create_infos[i];

        rse_free((VkPipelineVertexInputStateCreateInfo*)pipeline_info->pVertexInputState);
        rse_free((VkPipelineDynamicStateCreateInfo*)pipeline_info->pDynamicState);
    }

    pipelines->pipeline_create_infos_count = 0;
}

rse_err_t pipeline_add_dynamic_state(struct pipeline_t* pipeline, VkDynamicState dynamic_state)
{
    if (dynamic_state >= VK_DYNAMIC_STATE_MAX_ENUM) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Invalid dynamic state in pipeline");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    // TODO: This was "Push Front" when this was still a list. This might be relevant
    pipeline->dynamic_states[pipeline->dynamic_states_count++] = dynamic_state;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t pipeline_builder_init(struct pipeline_t* pipeline)
{
    rse_err_t status = RSE_ERROR_NO_ERROR;

    if (pipeline == NULL) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Provided NULL pointer");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    rse_memset(pipeline, 0, sizeof(struct pipeline_t));

    return status;
}

rse_err_t pipeline_add_shader_stage(struct graphics_context_t* context,
                                    struct pipeline_t* pipeline,
                                    const rse_id_t shader_id)
{
    size_t i = 0;
    rse_err_t status = RSE_ERROR_NO_ERROR;

    /* Check if we have enough space for new shader stages */
    if (pipeline->shader_stages_count + 1 > MAX_SHADER_STAGES) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Too many shader stages in pipeline");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    i = pipeline->shader_stages_count++;

    pipeline->shader_stages[i].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    pipeline->shader_stages[i].pNext = NULL;
    pipeline->shader_stages[i].flags = 0U;
    pipeline->shader_stages[i].stage = context->shader_modules.shaders_stages[shader_id];
    /* This is automatically updated during shader module creation */
    pipeline->shader_stages[i].module = context->shader_modules.shader_modules[shader_id];
    pipeline->shader_stages[i].pName = "main";
    pipeline->shader_stages[i].pSpecializationInfo = NULL;

    return status;
}

rse_err_t pipeline_add_vertex_input_binding(struct pipeline_t* pipeline,
                                            const uint8_t binding,
                                            const uint8_t stride,
                                            const VkVertexInputRate input_rate)
{
    VkVertexInputBindingDescription binding_pipeline_description = {0};
    if (input_rate >= VK_VERTEX_INPUT_RATE_MAX_ENUM) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Invalid vertex input rate");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    binding_pipeline_description.binding = binding;
    binding_pipeline_description.stride = stride;
    binding_pipeline_description.inputRate = input_rate;

    pipeline->vertex_bindings[pipeline->vertex_bindings_count++] = binding_pipeline_description;
    return RSE_ERROR_NO_ERROR;
}

rse_err_t pipeline_add_vertex_input_attribute(struct pipeline_t* pipeline,
                                              const uint8_t binding,
                                              const uint8_t location,
                                              const VkFormat format,
                                              const uint8_t offset)
{
    VkVertexInputAttributeDescription attribute_description = {0};
    if (binding >= 16U) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Invalid binding in pipeline attribute description");
        return RSE_ERROR_INTERNAL_ERROR;
    }
    if (format >= VK_FORMAT_MAX_ENUM) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Invalid format in pipeline attribute description");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    attribute_description.binding = binding;
    attribute_description.location = location;
    attribute_description.format = format;
    attribute_description.offset = offset;

    pipeline->shader_attributes[pipeline->shader_attributes_count++] = attribute_description;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t pipeline_add_push_constants(struct pipeline_t* pipeline,
                                       VkShaderStageFlagBits shader_stage,
                                       uint32_t offset,
                                       uint32_t size)
{
    pipeline->push_constant_ranges.push_constant_ranges[pipeline->push_constant_ranges.push_constant_ranges_count]
        .stageFlags = shader_stage;
    pipeline->push_constant_ranges.push_constant_ranges[pipeline->push_constant_ranges.push_constant_ranges_count]
        .offset = offset;
    pipeline->push_constant_ranges.push_constant_ranges[pipeline->push_constant_ranges.push_constant_ranges_count]
        .size = size;

    pipeline->push_constant_ranges.push_constant_ranges_count++;

    return RSE_ERROR_NO_ERROR;
}

rse_err_t pipelines_build(struct graphics_context_t* context, struct pipeline_infos_t* pipelines_infos_list)
{
    assert(NULL != context);
    assert(NULL != pipelines_infos_list);

    rse_err_t status = RSE_ERROR_NO_ERROR;

    if (VK_SUCCESS !=
        vkCreateGraphicsPipelines(context->vulkan_handles.device,
                                  VK_NULL_HANDLE,
                                  pipelines_infos_list->pipeline_create_infos_count,
                                  pipelines_infos_list->create_infos,
                                  NULL,
                                  &context->pipelines_data.pipelines[context->pipelines_data.pipelines_count])) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Failed to create graphics pipelines");
        status = RSE_ERROR_INTERNAL_ERROR;
    }

    context->pipelines_data.pipelines_count += pipelines_infos_list->pipeline_create_infos_count;
    pipeline_infos_cleanup(pipelines_infos_list);
    rse_memset(pipelines_infos_list, 0, sizeof(struct pipeline_infos_t));

    return status;
}

rse_err_t pipeline_add(const struct graphics_context_t* context,
                       struct pipeline_t* pipeline,
                       struct pipeline_infos_t* pipelines_infos)
{
    /* Check for send and required dynamic states */
    size_t iter = 0U;
    uint8_t dynamic_state_found = 0;
    rse_err_t status = RSE_ERROR_NO_ERROR;
    VkGraphicsPipelineCreateInfo pipeline_info = {0};
    VkPipelineVertexInputStateCreateInfo* vertex_bindings_ci = NULL;
    VkPipelineDynamicStateCreateInfo* dynamic_states_ci = NULL;

    while (pipeline->required_dynamic_states.dynamic_states_count > 0) {
        dynamic_state_found = 0;
        VkDynamicState dynamic_state =
            pipeline->required_dynamic_states.dynamic_states[pipeline->required_dynamic_states.dynamic_states_count--];
        for (iter = 0U; iter < pipeline->dynamic_states_count; ++iter) {
            if (pipeline->dynamic_states[iter] == dynamic_state) {
                dynamic_state_found = 1;
                break;
            }
        }
        if (dynamic_state_found == 0) {
            SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Missing dynamic state in pipeline");
            return RSE_ERROR_INTERNAL_ERROR;
        }
    }

    /* Check for required shader stages */
    if (pipeline->shader_stages_count == 0) {
        SDL_LogCritical(SDL_LOG_CATEGORY_GPU, "Missing shader stages in pipeline. There must be at least one.");
        return RSE_ERROR_INTERNAL_ERROR;
    }

    /* Build constant pipeline. TODO: Some parts may be customizable */
    pipeline_add_input_assembly(pipeline);
    pipeline_add_viewport_state(pipeline);
    pipeline_add_rasterizer(pipeline);
    pipeline_add_multisampling(pipeline);
    pipeline_add_color_blending(pipeline);
    pipeline_add_depth_stencil(pipeline);

    pipeline_info.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipeline_info.pNext = NULL;
    pipeline_info.flags = 0U;
    pipeline_info.stageCount = pipeline->shader_stages_count;
    pipeline_info.pStages = pipeline->shader_stages;

    rse_malloc(vertex_bindings_ci, sizeof(VkPipelineVertexInputStateCreateInfo));
    vertex_bindings_ci->sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertex_bindings_ci->pNext = NULL;
    vertex_bindings_ci->flags = 0U;
    vertex_bindings_ci->vertexBindingDescriptionCount = pipeline->vertex_bindings_count;
    vertex_bindings_ci->pVertexBindingDescriptions = pipeline->vertex_bindings;
    vertex_bindings_ci->vertexAttributeDescriptionCount = pipeline->shader_attributes_count;
    vertex_bindings_ci->pVertexAttributeDescriptions = pipeline->shader_attributes;
    pipeline_info.pVertexInputState = vertex_bindings_ci;

    pipeline_info.pInputAssemblyState = &pipeline->input_assembly;
    pipeline_info.pTessellationState = NULL;
    pipeline_info.pViewportState = &pipeline->viewport_state;
    pipeline_info.pRasterizationState = &pipeline->rasterizer;
    pipeline_info.pMultisampleState = &pipeline->multisampling;
    pipeline_info.pDepthStencilState = &pipeline->depth_stencil;
    pipeline_info.pColorBlendState = &pipeline->color_blending;

    rse_malloc(dynamic_states_ci, sizeof(VkPipelineDynamicStateCreateInfo));
    dynamic_states_ci->sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamic_states_ci->pNext = NULL;
    dynamic_states_ci->flags = 0U;
    dynamic_states_ci->dynamicStateCount = pipeline->dynamic_states_count;
    dynamic_states_ci->pDynamicStates = pipeline->dynamic_states;
    pipeline_info.pDynamicState = dynamic_states_ci;

    pipeline_info.layout = pipeline->pipeline_layout;
    pipeline_info.renderPass = context->vulkan_handles.render_pass;
    pipeline_info.subpass = 0;
    pipeline_info.basePipelineHandle = VK_NULL_HANDLE;  // TODO: Add possiblity do derive from another pipeline
    pipeline_info.basePipelineIndex = -1;

    pipelines_infos->create_infos[pipelines_infos->pipeline_create_infos_count++] = pipeline_info;

    return status;
}

void destroy_pipelines(struct graphics_context_t* context)
{
    size_t i = 0;

    for (i = 0; i < context->pipelines_data.pipelines_count; i++) {
        if (context->pipelines_data.pipeline_layouts[i] != VK_NULL_HANDLE) {
            vkDestroyPipelineLayout(context->vulkan_handles.device, context->pipelines_data.pipeline_layouts[i], NULL);
        }
        if (context->pipelines_data.pipelines[i] != VK_NULL_HANDLE) {
            vkDestroyPipeline(context->vulkan_handles.device, context->pipelines_data.pipelines[i], NULL);
        }
    }

    for (i = 0; i < context->shader_modules.shader_modules_count; i++) {
        vkDestroyShaderModule(context->vulkan_handles.device, context->shader_modules.shader_modules[i], NULL);
    }
}
