// =============================================
//  Aster: pipeline_utils.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "pipeline_utils.h"

#include "device.h"
#include "gpu_resource_manager.h"
#include "helpers.h"

#include <EASTL/array.h>

Pipeline
CreatePipeline(const Device *device, vk::Format attachmentFormat, const GpuResourceManager *resourceManager)
{
    // Pipeline Setup
    auto vertexShaderModule = CreateShader(device, VERTEX_SHADER_FILE);
    auto fragmentShaderModule = CreateShader(device, FRAGMENT_SHADER_FILE);

    eastl::array<vk::PipelineShaderStageCreateInfo, 2> shaderStages = {{
        {
            .stage = vk::ShaderStageFlagBits::eVertex,
            .module = vertexShaderModule,
            .pName = "main",
        },
        {
            .stage = vk::ShaderStageFlagBits::eFragment,
            .module = fragmentShaderModule,
            .pName = "main",
        },
    }};

    eastl::vector<vk::DescriptorSetLayout> descriptorSetLayouts;

    descriptorSetLayouts.push_back(resourceManager->m_SetLayout);

    {
        eastl::array descriptorSetLayoutBindings = {
            vk::DescriptorSetLayoutBinding{
                .binding = 0,
                .descriptorType = vk::DescriptorType::eUniformBuffer,
                .descriptorCount = 1,
                .stageFlags = vk::ShaderStageFlagBits::eAll,
            },
        };
        vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutCreateInfo = {
            .bindingCount = Cast<u32>(descriptorSetLayoutBindings.size()),
            .pBindings = descriptorSetLayoutBindings.data(),
        };
        vk::DescriptorSetLayout descriptorSetLayout;
        AbortIfFailed(
            device->m_Device.createDescriptorSetLayout(&descriptorSetLayoutCreateInfo, nullptr, &descriptorSetLayout));
        descriptorSetLayouts.push_back(descriptorSetLayout);
    }

    vk::PushConstantRange pushConstantRange = {
        .stageFlags = vk::ShaderStageFlagBits::eAll,
        .offset = 0,
        .size = 40,
    };

    vk::PipelineLayoutCreateInfo pipelineLayoutCreateInfo = {
        .setLayoutCount = Cast<u32>(descriptorSetLayouts.size()),
        .pSetLayouts = descriptorSetLayouts.data(),
        .pushConstantRangeCount = 1,
        .pPushConstantRanges = &pushConstantRange,
    };
    vk::PipelineLayout pipelineLayout;
    AbortIfFailed(device->m_Device.createPipelineLayout(&pipelineLayoutCreateInfo, nullptr, &pipelineLayout));
    device->SetName(pipelineLayout, "Box Layout");

    descriptorSetLayouts[0] = nullptr; // Not owned.

    vk::PipelineVertexInputStateCreateInfo vertexInputStateCreateInfo = {};
    vk::PipelineInputAssemblyStateCreateInfo inputAssemblyStateCreateInfo = {
        .topology = vk::PrimitiveTopology::eTriangleList,
        .primitiveRestartEnable = false,
    };

    vk::PipelineViewportStateCreateInfo viewportStateCreateInfo = {
        .viewportCount = 1,
        .scissorCount = 1,
    };

    vk::PipelineRasterizationStateCreateInfo rasterizationStateCreateInfo = {
        .depthClampEnable = false,
        .rasterizerDiscardEnable = false,
        .polygonMode = vk::PolygonMode::eFill,
        .cullMode = vk::CullModeFlagBits::eBack,
        .frontFace = vk::FrontFace::eCounterClockwise,
        .depthBiasEnable = false,
        .lineWidth = 1.0,
    };
    vk::PipelineMultisampleStateCreateInfo multisampleStateCreateInfo = {
        .rasterizationSamples = vk::SampleCountFlagBits::e1,
        .sampleShadingEnable = false,
    };
    vk::PipelineDepthStencilStateCreateInfo depthStencilStateCreateInfo = {
        .depthTestEnable = true,
        .depthWriteEnable = true,
        .depthCompareOp = vk::CompareOp::eLess,
    };
    vk::PipelineColorBlendAttachmentState colorBlendAttachmentState = {
        .blendEnable = false,
        .srcColorBlendFactor = vk::BlendFactor::eSrcColor,
        .dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcColor,
        .colorBlendOp = vk::BlendOp::eAdd,
        .srcAlphaBlendFactor = vk::BlendFactor::eSrcAlpha,
        .dstAlphaBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha,
        .alphaBlendOp = vk::BlendOp::eAdd,
        .colorWriteMask = vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
                          vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA,
    };
    vk::PipelineColorBlendStateCreateInfo colorBlendStateCreateInfo = {
        .logicOpEnable = false,
        .attachmentCount = 1,
        .pAttachments = &colorBlendAttachmentState,
    };

    eastl::array dynamicStates = {
        vk::DynamicState::eScissor,
        vk::DynamicState::eViewport,
    };

    vk::PipelineDynamicStateCreateInfo dynamicStateCreateInfo = {
        .dynamicStateCount = Cast<u32>(dynamicStates.size()),
        .pDynamicStates = dynamicStates.data(),
    };

    vk::PipelineRenderingCreateInfo renderingCreateInfo = {
        .viewMask = 0,
        .colorAttachmentCount = 1,
        .pColorAttachmentFormats = &attachmentFormat,
        .depthAttachmentFormat = vk::Format::eD24UnormS8Uint,
    };

    vk::GraphicsPipelineCreateInfo pipelineCreateInfo = {
        .pNext = &renderingCreateInfo,
        .stageCount = Cast<u32>(shaderStages.size()),
        .pStages = shaderStages.data(),
        .pVertexInputState = &vertexInputStateCreateInfo,
        .pInputAssemblyState = &inputAssemblyStateCreateInfo,
        .pViewportState = &viewportStateCreateInfo,
        .pRasterizationState = &rasterizationStateCreateInfo,
        .pMultisampleState = &multisampleStateCreateInfo,
        .pDepthStencilState = &depthStencilStateCreateInfo,
        .pColorBlendState = &colorBlendStateCreateInfo,
        .pDynamicState = &dynamicStateCreateInfo,
        .layout = pipelineLayout,
    };
    vk::Pipeline pipeline;
    AbortIfFailed(
        device->m_Device.createGraphicsPipelines(device->m_PipelineCache, 1, &pipelineCreateInfo, nullptr, &pipeline));
    device->SetName(pipeline, "Box Pipeline");

    device->m_Device.destroy(vertexShaderModule, nullptr);
    device->m_Device.destroy(fragmentShaderModule, nullptr);

    return {device, pipelineLayout, pipeline, std::move(descriptorSetLayouts)};
}

Pipeline
CreateBackgroundPipeline(const Device *device, vk::Format attachmentFormat, const GpuResourceManager *resourceManager)
{
    // Pipeline Setup
    auto vertexShaderModule = CreateShader(device, BACKGROUND_VERTEX_SHADER_FILE);
    auto fragmentShaderModule = CreateShader(device, BACKGROUND_FRAGMENT_SHADER_FILE);

    eastl::array<vk::PipelineShaderStageCreateInfo, 2> shaderStages = {{
        {
            .stage = vk::ShaderStageFlagBits::eVertex,
            .module = vertexShaderModule,
            .pName = "main",
        },
        {
            .stage = vk::ShaderStageFlagBits::eFragment,
            .module = fragmentShaderModule,
            .pName = "main",
        },
    }};

    eastl::vector<vk::DescriptorSetLayout> descriptorSetLayouts;

    descriptorSetLayouts.push_back(resourceManager->m_SetLayout);

    {
        eastl::array descriptorSetLayoutBindings = {
            vk::DescriptorSetLayoutBinding{
                .binding = 0,
                .descriptorType = vk::DescriptorType::eUniformBuffer,
                .descriptorCount = 1,
                .stageFlags = vk::ShaderStageFlagBits::eAll,
            },
        };
        vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutCreateInfo = {
            .bindingCount = Cast<u32>(descriptorSetLayoutBindings.size()),
            .pBindings = descriptorSetLayoutBindings.data(),
        };
        vk::DescriptorSetLayout descriptorSetLayout;
        AbortIfFailed(
            device->m_Device.createDescriptorSetLayout(&descriptorSetLayoutCreateInfo, nullptr, &descriptorSetLayout));
        descriptorSetLayouts.push_back(descriptorSetLayout);
    }

    vk::PushConstantRange pushConstantRange = {
        .stageFlags = vk::ShaderStageFlagBits::eAll,
        .offset = 0,
        .size = 40,
    };

    vk::PipelineLayoutCreateInfo pipelineLayoutCreateInfo = {
        .setLayoutCount = Cast<u32>(descriptorSetLayouts.size()),
        .pSetLayouts = descriptorSetLayouts.data(),
        .pushConstantRangeCount = 1,
        .pPushConstantRanges = &pushConstantRange,
    };
    vk::PipelineLayout pipelineLayout;
    AbortIfFailed(device->m_Device.createPipelineLayout(&pipelineLayoutCreateInfo, nullptr, &pipelineLayout));
    device->SetName(pipelineLayout, "Box Layout");

    descriptorSetLayouts[0] = nullptr; // Not owned.

    vk::PipelineVertexInputStateCreateInfo vertexInputStateCreateInfo = {};
    vk::PipelineInputAssemblyStateCreateInfo inputAssemblyStateCreateInfo = {
        .topology = vk::PrimitiveTopology::eTriangleList,
        .primitiveRestartEnable = false,
    };

    vk::PipelineViewportStateCreateInfo viewportStateCreateInfo = {
        .viewportCount = 1,
        .scissorCount = 1,
    };

    vk::PipelineRasterizationStateCreateInfo rasterizationStateCreateInfo = {
        .depthClampEnable = false,
        .rasterizerDiscardEnable = false,
        .polygonMode = vk::PolygonMode::eFill,
        .cullMode = vk::CullModeFlagBits::eBack,
        .frontFace = vk::FrontFace::eCounterClockwise,
        .depthBiasEnable = false,
        .lineWidth = 1.0,
    };
    vk::PipelineMultisampleStateCreateInfo multisampleStateCreateInfo = {
        .rasterizationSamples = vk::SampleCountFlagBits::e1,
        .sampleShadingEnable = false,
    };
    vk::PipelineDepthStencilStateCreateInfo depthStencilStateCreateInfo = {
        .depthTestEnable = true,
        .depthWriteEnable = true,
        .depthCompareOp = vk::CompareOp::eLessOrEqual,
    };
    vk::PipelineColorBlendAttachmentState colorBlendAttachmentState = {
        .blendEnable = false,
        .srcColorBlendFactor = vk::BlendFactor::eSrcColor,
        .dstColorBlendFactor = vk::BlendFactor::eOneMinusSrcColor,
        .colorBlendOp = vk::BlendOp::eAdd,
        .srcAlphaBlendFactor = vk::BlendFactor::eSrcAlpha,
        .dstAlphaBlendFactor = vk::BlendFactor::eOneMinusSrcAlpha,
        .alphaBlendOp = vk::BlendOp::eAdd,
        .colorWriteMask = vk::ColorComponentFlagBits::eR | vk::ColorComponentFlagBits::eG |
                          vk::ColorComponentFlagBits::eB | vk::ColorComponentFlagBits::eA,
    };
    vk::PipelineColorBlendStateCreateInfo colorBlendStateCreateInfo = {
        .logicOpEnable = false,
        .attachmentCount = 1,
        .pAttachments = &colorBlendAttachmentState,
    };

    eastl::array dynamicStates = {
        vk::DynamicState::eScissor,
        vk::DynamicState::eViewport,
    };

    vk::PipelineDynamicStateCreateInfo dynamicStateCreateInfo = {
        .dynamicStateCount = Cast<u32>(dynamicStates.size()),
        .pDynamicStates = dynamicStates.data(),
    };

    vk::PipelineRenderingCreateInfo renderingCreateInfo = {
        .viewMask = 0,
        .colorAttachmentCount = 1,
        .pColorAttachmentFormats = &attachmentFormat,
        .depthAttachmentFormat = vk::Format::eD24UnormS8Uint,
    };

    vk::GraphicsPipelineCreateInfo pipelineCreateInfo = {
        .pNext = &renderingCreateInfo,
        .stageCount = Cast<u32>(shaderStages.size()),
        .pStages = shaderStages.data(),
        .pVertexInputState = &vertexInputStateCreateInfo,
        .pInputAssemblyState = &inputAssemblyStateCreateInfo,
        .pViewportState = &viewportStateCreateInfo,
        .pRasterizationState = &rasterizationStateCreateInfo,
        .pMultisampleState = &multisampleStateCreateInfo,
        .pDepthStencilState = &depthStencilStateCreateInfo,
        .pColorBlendState = &colorBlendStateCreateInfo,
        .pDynamicState = &dynamicStateCreateInfo,
        .layout = pipelineLayout,
    };
    vk::Pipeline pipeline;
    AbortIfFailed(
        device->m_Device.createGraphicsPipelines(device->m_PipelineCache, 1, &pipelineCreateInfo, nullptr, &pipeline));
    device->SetName(pipeline, "BG Pipeline");

    device->m_Device.destroy(vertexShaderModule, nullptr);
    device->m_Device.destroy(fragmentShaderModule, nullptr);

    return {device, pipelineLayout, pipeline, std::move(descriptorSetLayouts)};
}

vk::ShaderModule
CreateShader(const Device *device, cstr shaderFile)
{
    eastl::vector<u32> shaderCode = ReadFile(shaderFile);

    const vk::ShaderModuleCreateInfo shaderModuleCreateInfo = {
        .codeSize = shaderCode.size() * sizeof(u32),
        .pCode = shaderCode.data(),
    };
    vk::ShaderModule shaderModule;

    AbortIfFailedMV(device->m_Device.createShaderModule(&shaderModuleCreateInfo, nullptr, &shaderModule),
                    "Shader {} could not be created.", shaderFile);
    return shaderModule;
}