#include "iwa/pipeline.hpp"

#include <mijin/io/stream.hpp>
#include <yaml-cpp/yaml.h>
#include "iwa/device.hpp"
#include "iwa/util/glsl_compiler.hpp"
#include "iwa/util/next_chain.hpp"

namespace iwa
{
PipelineLayout::PipelineLayout(ObjectPtr<Device> owner, const PipelineLayoutCreationArgs& args) : super_t(std::move(owner))
{
    std::vector<vk::DescriptorSetLayout> setLayoutHandles;
    setLayoutHandles.reserve(args.setLayouts.size());
    for (const ObjectPtr<DescriptorSetLayout>& setLayout : args.setLayouts)
    {
        setLayoutHandles.push_back(setLayout->getVkHandle());
    }

    mHandle = getOwner()->getVkHandle().createPipelineLayout(vk::PipelineLayoutCreateInfo{
        .flags = args.flags,
        .setLayoutCount = static_cast<std::uint32_t>(setLayoutHandles.size()),
        .pSetLayouts = setLayoutHandles.data(),
        .pushConstantRangeCount = static_cast<std::uint32_t>(args.pushConstantRanges.size()),
        .pPushConstantRanges = args.pushConstantRanges.data()
    });
}

PipelineLayout::~PipelineLayout() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyPipelineLayout)
}

Pipeline::Pipeline(ObjectPtr<Device> owner) noexcept : super_t(std::move(owner))
{
}

Pipeline::~Pipeline() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyPipeline)
}

GraphicsPipeline::GraphicsPipeline(ObjectPtr<Device> owner, const GraphicsPipelineCreationArgs& args) noexcept
        : super_t(std::move(owner))
{
    std::vector<vk::PipelineShaderStageCreateInfo> vkStages;

    for (const PipelineStage& stage : args.stages)
    {
        vk::PipelineShaderStageCreateInfo& vkStage = vkStages.emplace_back();
        vkStage.stage = stage.stage;
        vkStage.module = *stage.shader;
        vkStage.pName = stage.name.c_str();
    }

    const vk::PipelineVertexInputStateCreateInfo vkVertexInput =
    {
        .vertexBindingDescriptionCount   = static_cast<std::uint32_t>(args.vertexInput.bindings.size()),
        .pVertexBindingDescriptions      = args.vertexInput.bindings.data(),
        .vertexAttributeDescriptionCount = static_cast<std::uint32_t>(args.vertexInput.attributes.size()),
        .pVertexAttributeDescriptions    = args.vertexInput.attributes.data()
    };

    const vk::PipelineColorBlendStateCreateInfo vkColorBlend =
    {
        .logicOpEnable = args.colorBlend.logicOp.has_value(),
        .logicOp = args.colorBlend.logicOp.has_value() ? *args.colorBlend.logicOp : vk::LogicOp(),
        .attachmentCount = static_cast<std::uint32_t>(args.colorBlend.attachements.size()),
        .pAttachments = args.colorBlend.attachements.data()
    };

    const vk::PipelineDynamicStateCreateInfo vkDynamic =
    {
        .dynamicStateCount = static_cast<std::uint32_t>(args.dynamicState.size()),
        .pDynamicStates = args.dynamicState.data()
    };

    NextChain nextChain;
    if (args.renderingInfo.has_value())
    {
        const GraphicsPipelineRenderingInfo& rinfo = *args.renderingInfo;
        nextChain.append(vk::PipelineRenderingCreateInfo{
            .viewMask = rinfo.viewMask,
            .colorAttachmentCount = static_cast<std::uint32_t>(rinfo.colorAttachmentFormats.size()),
            .pColorAttachmentFormats = rinfo.colorAttachmentFormats.data(),
            .depthAttachmentFormat = rinfo.depthFormat,
            .stencilAttachmentFormat = rinfo.stencilFormat
        });
    }

    const vk::GraphicsPipelineCreateInfo vkCreateInfo =
    {
        .pNext               = nextChain.finalize(),
        .stageCount          = static_cast<std::uint32_t>(vkStages.size()),
        .pStages             = vkStages.data(),
        .pVertexInputState   = &vkVertexInput,
        .pInputAssemblyState = &args.inputAssembly,
        .pViewportState      = &args.viewport,
        .pRasterizationState = &args.rasterization,
        .pMultisampleState   = &args.multisample,
        .pDepthStencilState  = &args.depthStencil,
        .pColorBlendState    = &vkColorBlend,
        .pDynamicState       = &vkDynamic,
        .layout              = *args.layout,
        .renderPass          = args.renderPass ? *args.renderPass : vk::RenderPass(),
        .subpass             = args.subpass
    };

    const auto result = getOwner()->getVkHandle().createGraphicsPipeline(VK_NULL_HANDLE, vkCreateInfo);
    MIJIN_ASSERT(result.result == vk::Result::eSuccess, "Graphics pipeline creation failed.");
    mHandle = result.value;
}

ComputePipeline::ComputePipeline(ObjectPtr<Device> owner, const ComputePipelineCreationArgs& args) noexcept
        : super_t(std::move(owner))
{
    const vk::ComputePipelineCreateInfo vkCreateInfo{
        .stage = {
            .stage = vk::ShaderStageFlagBits::eCompute,
            .module = *args.stage.shader,
            .pName = args.stage.name.c_str()
        },
        .layout = *args.layout
    };
    const vk::ResultValue<vk::Pipeline> result = getOwner()->getVkHandle().createComputePipeline(VK_NULL_HANDLE, vkCreateInfo);
    MIJIN_ASSERT(result.result == vk::Result::eSuccess, "Compute pipeline creation failed.");
    mHandle = result.value;
}

RayTracingPipeline::RayTracingPipeline(ObjectPtr<Device> owner, const RayTracingPipelineCreationArgs& args) noexcept
        : super_t(std::move(owner))
{
    (void) args;
}
} // namespace iwa