#include <Vk/Pipeline/Pipeline.hpp>
#include <GraphicsContext/GraphicsContext.hpp>
#include <Utilities/Utilities.hpp>

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addShaderStage(VkShaderStageFlagBits stage, Vk::ShaderModule &shaderModule)
{
    VkPipelineShaderStageCreateInfo shaderStageCreateInfo{};
    shaderStageCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    shaderStageCreateInfo.stage = stage;
    shaderStageCreateInfo.module = shaderModule.inner();
    shaderStageCreateInfo.pName = "main";

    shaderStages.push_back(shaderStageCreateInfo);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addDynamicState(VkDynamicState state)
{
    dynamicStates.push_back(state);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addPushConstantRange(
    uint32_t offset,
    uint32_t size,
    VkShaderStageFlags stages)
{
    VkPushConstantRange range{};
    range.stageFlags = stages;
    range.offset = offset;
    range.size = size;
    pushConstantRanges.push_back(range);
    
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addDescriptorSetLayout(std::shared_ptr<Vk::DescriptorSetLayout> layout)
{
    descriptorSetLayoutHandles.push_back(layout);
    descriptorSetLayouts.push_back(layout->inner());

    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::enableDepthTest()
{
    m_depthTest = true;
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::enableDepthWrite()
{
    m_depthWrite = true;
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addAlphaBlendedAttachment()
{
    VkPipelineColorBlendAttachmentState colorBlendAttachment{};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

    blendModes.push_back(colorBlendAttachment);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addAdditiveBlendedAttachment()
{
    VkPipelineColorBlendAttachmentState colorBlendAttachment{};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

    blendModes.push_back(colorBlendAttachment);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addInvertBlendedAttachment()
{
    VkPipelineColorBlendAttachmentState colorBlendAttachment{};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = VK_TRUE;
    colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_COLOR;
    colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
    colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

    blendModes.push_back(colorBlendAttachment);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addNonBlendedAttachment()
{
    VkPipelineColorBlendAttachmentState colorBlendAttachment{};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = VK_FALSE;

    blendModes.push_back(colorBlendAttachment);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addVertexInputBinding(int binding, size_t stride, VkVertexInputRate inputRate)
{
    VkVertexInputBindingDescription bindingDescription{};
    bindingDescription.binding = binding;
    bindingDescription.stride = stride;
    bindingDescription.inputRate = inputRate;

    vertexInputBindings.push_back(bindingDescription);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::addVertexAttr(int binding, int location, VkFormat format, size_t offset)
{
    VkVertexInputAttributeDescription attr;
    attr.binding = binding;
    attr.location = location;
    attr.format = format;
    attr.offset = offset;

    vertexAttrs.push_back(attr);
    return *this;
}

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::setFrontFace(VkFrontFace frontFace)
{
    m_frontFace = frontFace;
    return *this;
};

Vk::Pipeline::GraphicsBuilder& Vk::Pipeline::GraphicsBuilder::setCullMode(VkCullModeFlags cullMode)
{
    m_cullMode = cullMode;
    return *this;
}

std::shared_ptr<Vk::Pipeline> Vk::Pipeline::GraphicsBuilder::build(
    std::shared_ptr<Vk::Device> device,
    std::shared_ptr<Vk::RenderPass> renderPass
)
{
    VkPipelineDynamicStateCreateInfo dynamicStateCreateInfo{};
    dynamicStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicStateCreateInfo.dynamicStateCount = dynamicStates.size();
    dynamicStateCreateInfo.pDynamicStates = dynamicStates.data();

    VkPipelineVertexInputStateCreateInfo vertexInputCreateInfo{};
    vertexInputCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputCreateInfo.vertexBindingDescriptionCount = vertexInputBindings.size();
    vertexInputCreateInfo.pVertexBindingDescriptions = vertexInputBindings.data();
    vertexInputCreateInfo.vertexAttributeDescriptionCount = vertexAttrs.size();
    vertexInputCreateInfo.pVertexAttributeDescriptions = vertexAttrs.data();

    VkPipelineInputAssemblyStateCreateInfo inputAssemblyStateCreateInfo{};
    inputAssemblyStateCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssemblyStateCreateInfo.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    inputAssemblyStateCreateInfo.primitiveRestartEnable = VK_FALSE;

    VkPipelineViewportStateCreateInfo viewportState{};
    viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportState.viewportCount = 1;
    viewportState.scissorCount = 1;

    VkPipelineRasterizationStateCreateInfo rasterizer{};
    rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterizer.depthClampEnable = VK_FALSE;
    rasterizer.rasterizerDiscardEnable = VK_FALSE;
    rasterizer.polygonMode = VK_POLYGON_MODE_FILL;
    rasterizer.lineWidth = 1.0f;
    rasterizer.cullMode = m_cullMode;
    rasterizer.frontFace = m_frontFace;
    rasterizer.depthBiasEnable = VK_FALSE;
    rasterizer.depthBiasConstantFactor = 0.0f;
    rasterizer.depthBiasClamp = 0.0f;
    rasterizer.depthBiasSlopeFactor = 0.0f;

    VkPipelineMultisampleStateCreateInfo multisampling{};
    multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisampling.sampleShadingEnable = VK_FALSE;
    multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
    multisampling.minSampleShading = 1.0f;
    multisampling.pSampleMask = nullptr;
    multisampling.alphaToCoverageEnable = VK_FALSE;
    multisampling.alphaToOneEnable = VK_FALSE;

    VkPipelineColorBlendStateCreateInfo colorBlending{};
    colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlending.logicOpEnable = VK_FALSE;
    colorBlending.logicOp = VK_LOGIC_OP_COPY;
    colorBlending.attachmentCount = blendModes.size();
    colorBlending.pAttachments = blendModes.data();
    colorBlending.blendConstants[0] = 0.0f;
    colorBlending.blendConstants[1] = 0.0f;
    colorBlending.blendConstants[2] = 0.0f;
    colorBlending.blendConstants[3] = 0.0f;

    VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = descriptorSetLayouts.size();
    pipelineLayoutInfo.pSetLayouts = descriptorSetLayouts.data();
    pipelineLayoutInfo.pushConstantRangeCount = pushConstantRanges.size();
    pipelineLayoutInfo.pPushConstantRanges = pushConstantRanges.data();

    VkPipelineLayout pipelineLayout;
    if (vkCreatePipelineLayout(device->inner(), &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS)
    {
        crash("Failed to create pipeline layout!");
    }

    VkPipelineDepthStencilStateCreateInfo depthStencil{};
    depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    depthStencil.depthTestEnable = (m_depthTest || m_depthWrite) ? VK_TRUE : VK_FALSE;
    depthStencil.depthWriteEnable = m_depthWrite ? VK_TRUE : VK_FALSE;
    depthStencil.depthCompareOp = m_depthTest ? VK_COMPARE_OP_LESS : VK_COMPARE_OP_ALWAYS;

    VkGraphicsPipelineCreateInfo pipelineInfo{};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = shaderStages.size();
    pipelineInfo.pStages = shaderStages.data();
    pipelineInfo.pVertexInputState = &vertexInputCreateInfo;
    pipelineInfo.pInputAssemblyState = &inputAssemblyStateCreateInfo;
    pipelineInfo.pViewportState = &viewportState;
    pipelineInfo.pRasterizationState = &rasterizer;
    pipelineInfo.pMultisampleState = &multisampling;
    pipelineInfo.pDepthStencilState = nullptr;
    pipelineInfo.pColorBlendState = &colorBlending;
    pipelineInfo.pDynamicState = &dynamicStateCreateInfo;
    pipelineInfo.layout = pipelineLayout;
    pipelineInfo.renderPass = renderPass->inner();
    pipelineInfo.subpass = 0;
    pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;
    pipelineInfo.basePipelineIndex = -1;

    if (m_depthTest || m_depthWrite)
    {
        pipelineInfo.pDepthStencilState = &depthStencil;
    }

    VkPipeline pipeline;
    if (vkCreateGraphicsPipelines(device->inner(), VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &pipeline) != VK_SUCCESS)
    {
        crash("Failed to create graphics pipeline!");
    }

    return std::make_shared<Vk::Pipeline>(device, pipeline, pipelineLayout, renderPass);
}

Vk::Pipeline::ComputeBuilder& Vk::Pipeline::ComputeBuilder::addDescriptorSetLayout(std::shared_ptr<Vk::DescriptorSetLayout> layout)
{
    descriptorSetLayoutHandles.push_back(layout);
    descriptorSetLayouts.push_back(layout->inner());

    return *this;
}

Vk::Pipeline::ComputeBuilder& Vk::Pipeline::ComputeBuilder::addPushConstantRange(
    uint32_t offset,
    uint32_t size,
    VkShaderStageFlags stages)
{
    VkPushConstantRange range{};
    range.stageFlags = stages;
    range.offset = offset;
    range.size = size;
    pushConstantRanges.push_back(range);
    
    return *this;
}

std::shared_ptr<Vk::Pipeline> Vk::Pipeline::ComputeBuilder::build(std::shared_ptr<Device> device, Vk::ShaderModule &shaderModule)
{
    VkPipelineLayoutCreateInfo pipelineLayoutCreateInfo = {};
    pipelineLayoutCreateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutCreateInfo.setLayoutCount = descriptorSetLayouts.size();
    pipelineLayoutCreateInfo.pSetLayouts = descriptorSetLayouts.data();
    pipelineLayoutCreateInfo.pushConstantRangeCount = pushConstantRanges.size();
    pipelineLayoutCreateInfo.pPushConstantRanges = pushConstantRanges.data();

    VkPipelineLayout pipelineLayout;
    if (vkCreatePipelineLayout(device->inner(), &pipelineLayoutCreateInfo, nullptr, &pipelineLayout) != VK_SUCCESS)
    {
        crash("Failed to create pipeline layout!");
    }

    VkComputePipelineCreateInfo pipelineCreateInfo = {};
    pipelineCreateInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
    pipelineCreateInfo.stage = {
        VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO, 
        nullptr, 
        0, 
        VK_SHADER_STAGE_COMPUTE_BIT, 
        shaderModule.inner(), 
        "main"
    };
    pipelineCreateInfo.layout = pipelineLayout;

    VkPipeline pipeline;
    if (vkCreateComputePipelines(device->inner(), VK_NULL_HANDLE, 1, &pipelineCreateInfo, nullptr, &pipeline) != VK_SUCCESS)
    {
        crash("Failed to create compute pipeline!");
    }

    return std::make_shared<Vk::Pipeline>(device, pipeline, pipelineLayout, nullptr);
}

Vk::Pipeline::Pipeline(
    std::shared_ptr<Vk::Device> device,
    VkPipeline pipeline,
    VkPipelineLayout layout,
    std::shared_ptr<Vk::RenderPass> renderPass
) :
    device(device),
    pipeline(pipeline),
    layout(layout),
    renderPass(renderPass)
{}

Vk::Pipeline::~Pipeline()
{
    vkDestroyPipeline(device->inner(), pipeline, nullptr);
    vkDestroyPipelineLayout(device->inner(), layout, nullptr);
}

VkPipeline Vk::Pipeline::inner()
{
    return pipeline;
}

VkPipelineLayout Vk::Pipeline::getInnerLayout()
{
    return layout;
}