#include "Pipeline.hpp"

#include <array>
#include <stdexcept>
#include <fstream>

#include "DescriptorSetLayout.hpp"
#include "Device.hpp"
#include "RenderPass.hpp"
#include "Swapchain.hpp"
#include "DeviceFunctions/DeviceFunctions.hpp"
#include "../ModelManager.hpp"

VkShaderModule VkObjects::Pipeline::create_shader_module(const Device& device, const std::string& code) const
{
    VkShaderModuleCreateInfo shader_module_info{};
    shader_module_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    shader_module_info.codeSize = code.size();
    shader_module_info.pCode = reinterpret_cast<const uint32_t*>(code.c_str());

    VkShaderModule shader_module;
    VK_CHECK(vkCreateShaderModule(device.get_native(), &shader_module_info, nullptr, &shader_module))
    return shader_module;
}

VkObjects::Pipeline::Pipeline(const Device& device, const Swapchain& swapchain,
    const RenderPass& render_pass, const DescriptorSetLayout& layout_binding):device_(device)
{
    std::ifstream vertex_shader_file("Shaders/vertex.spv", std::ios::binary);
    if (!vertex_shader_file.is_open())
        throw std::runtime_error("Vertex shader not found");
    std::ifstream fragment_shader_file("Shaders/fragment.spv", std::ios::binary);
    if (!fragment_shader_file.is_open()) {
        vertex_shader_file.close();
        throw std::runtime_error("Fragment shader not found");
    }
    
    std::string vertex_shader_code((std::istreambuf_iterator<char>(vertex_shader_file)), std::istreambuf_iterator<char>());
    std::string fragment_shader_code((std::istreambuf_iterator<char>(fragment_shader_file)), std::istreambuf_iterator<char>());
    vertex_shader_file.close();
    fragment_shader_file.close();

    VkShaderModule vertex_shader = create_shader_module(device_, vertex_shader_code);
    VkShaderModule fragment_shader = create_shader_module(device_, fragment_shader_code);

    VkPipelineShaderStageCreateInfo vertex_shader_stage_info{};
    vertex_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    vertex_shader_stage_info.stage = VK_SHADER_STAGE_VERTEX_BIT;
    vertex_shader_stage_info.module = vertex_shader;
    vertex_shader_stage_info.pName = "main";

    VkPipelineShaderStageCreateInfo fragment_shader_stage_info{};
    fragment_shader_stage_info.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    fragment_shader_stage_info.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    fragment_shader_stage_info.module = fragment_shader;
    fragment_shader_stage_info.pName = "main";

    VkPipelineShaderStageCreateInfo shader_stages[] = {vertex_shader_stage_info, fragment_shader_stage_info};
    const std::vector<VkDynamicState> dynamic_states = {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR
    };

    VkPipelineDynamicStateCreateInfo dynamic_state_info{};
    dynamic_state_info.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamic_state_info.dynamicStateCount = static_cast<uint32_t>(dynamic_states.size());
    dynamic_state_info.pDynamicStates = dynamic_states.data();

    VkVertexInputBindingDescription binding_description = ModelManager::Vertex::get_binding_description();
    std::array<VkVertexInputAttributeDescription, 2> attribute_descriptions = ModelManager::Vertex::get_vertex_attribute_descriptions();
    
    VkPipelineVertexInputStateCreateInfo vertex_input_info{};
    vertex_input_info.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertex_input_info.vertexBindingDescriptionCount = 1;
    vertex_input_info.pVertexBindingDescriptions = &binding_description;
    vertex_input_info.vertexAttributeDescriptionCount = static_cast<uint32_t>(attribute_descriptions.size());
    vertex_input_info.pVertexAttributeDescriptions = attribute_descriptions.data();

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

    VkViewport viewport{};
    viewport.x = 0.0f;
    viewport.y = 0.0f;
    viewport.width = static_cast<float>(swapchain.get_extent().width);
    viewport.height = static_cast<float>(swapchain.get_extent().height);
    viewport.minDepth = 0.0f;
    viewport.maxDepth = 1.0f;

    VkRect2D scissor{};
    scissor.offset = {0, 0};
    scissor.extent = swapchain.get_extent();

    VkPipelineViewportStateCreateInfo viewport_info{};
    viewport_info.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewport_info.viewportCount = 1;
    viewport_info.pViewports = &viewport;
    viewport_info.scissorCount = 1;
    viewport_info.pScissors = &scissor;

    VkPipelineRasterizationStateCreateInfo rasterization_info{};
    rasterization_info.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterization_info.depthClampEnable = VK_FALSE;
    rasterization_info.rasterizerDiscardEnable = VK_FALSE;
    rasterization_info.polygonMode = VK_POLYGON_MODE_FILL;
    rasterization_info.lineWidth = 1.0f;
    rasterization_info.cullMode = VK_CULL_MODE_BACK_BIT;
    rasterization_info.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
    rasterization_info.depthBiasEnable = VK_FALSE;

    VkPipelineMultisampleStateCreateInfo multisample_info{};
    multisample_info.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisample_info.sampleShadingEnable = VK_FALSE;
    multisample_info.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;

    VkPipelineColorBlendAttachmentState color_blend_attachment_state{};
    color_blend_attachment_state.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
    color_blend_attachment_state.blendEnable = VK_TRUE;
    color_blend_attachment_state.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
    color_blend_attachment_state.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
    color_blend_attachment_state.colorBlendOp = VK_BLEND_OP_ADD;
    color_blend_attachment_state.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
    color_blend_attachment_state.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    color_blend_attachment_state.alphaBlendOp = VK_BLEND_OP_ADD;

    VkPipelineColorBlendStateCreateInfo color_blend_info{};
    color_blend_info.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    color_blend_info.logicOpEnable = VK_FALSE;
    color_blend_info.attachmentCount = 1;
    color_blend_info.pAttachments = &color_blend_attachment_state;
    color_blend_info.blendConstants[0] = 0.0f;
    color_blend_info.blendConstants[1] = 0.0f;
    color_blend_info.blendConstants[2] = 0.0f;
    color_blend_info.blendConstants[3] = 0.0f;

    VkDescriptorSetLayout layout = layout_binding.get_native();
    VkPipelineLayoutCreateInfo pipeline_layout_info{};
    pipeline_layout_info.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipeline_layout_info.setLayoutCount = 1;
    pipeline_layout_info.pSetLayouts = &layout;

    VK_CHECK(vkCreatePipelineLayout(device_.get_native(), &pipeline_layout_info, nullptr, &pipeline_layout_))

    VkGraphicsPipelineCreateInfo pipeline_info{};
    pipeline_info.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipeline_info.stageCount = 2;
    pipeline_info.pStages = shader_stages;
    pipeline_info.layout = pipeline_layout_;
    pipeline_info.pColorBlendState = &color_blend_info;
    pipeline_info.pDynamicState = &dynamic_state_info;
    pipeline_info.pInputAssemblyState = &input_assembly_info;
    pipeline_info.pMultisampleState = &multisample_info;
    pipeline_info.pRasterizationState = &rasterization_info;
    pipeline_info.pViewportState = &viewport_info;
    pipeline_info.pVertexInputState = &vertex_input_info;
    pipeline_info.pDepthStencilState = nullptr;
    pipeline_info.renderPass = render_pass.get_native();
    pipeline_info.subpass = 0;
    pipeline_info.basePipelineHandle = nullptr;
    pipeline_info.basePipelineIndex = -1;

    VK_CHECK(vkCreateGraphicsPipelines(device_.get_native(), nullptr, 1, &pipeline_info, nullptr, &pipeline_))
    
    vkDestroyShaderModule(device_.get_native(), vertex_shader, nullptr);
    vkDestroyShaderModule(device_.get_native(), fragment_shader, nullptr);
}

VkObjects::Pipeline::~Pipeline()
{
    if (pipeline_ != nullptr)
        vkDestroyPipeline(device_.get_native(), pipeline_, nullptr);
    
    if (pipeline_layout_ != nullptr)
        vkDestroyPipelineLayout(device_.get_native(), pipeline_layout_, nullptr);
}