#include <destrum/Graphics/Pipeline.h>

#include <cassert>
#include <filesystem>
#include <fstream>
#include <iostream>

#include "volk.h"
#include "destrum/Graphics/Util.h"
#include "spdlog/spdlog.h"


Pipeline::Pipeline(GfxDevice& device, const std::string& vertPath, const std::string& fragPath,
                   const PipelineConfigInfo& configInfo): m_device(device) {
    try {
        CreateGraphicsPipeline(vertPath, fragPath, configInfo);
    } catch (...) {
        if (m_device.getDevice() != VK_NULL_HANDLE) {
            vkDestroyShaderModule(m_device.getDevice(), m_vertShaderModule, nullptr);
            vkDestroyShaderModule(m_device.getDevice(), m_fragShaderModule, nullptr);
            vkDestroyPipeline(m_device.getDevice(), m_graphicsPipeline, nullptr);
        }
        throw;
    }
}

Pipeline::~Pipeline() {
    if (m_device.getDevice() != VK_NULL_HANDLE) {
        vkDestroyShaderModule(m_device.getDevice(), m_vertShaderModule, nullptr);
        vkDestroyShaderModule(m_device.getDevice(), m_fragShaderModule, nullptr);
        vkDestroyPipeline(m_device.getDevice(), m_graphicsPipeline, nullptr);
    }
}

void Pipeline::bind(VkCommandBuffer buffer) const {
    vkCmdBindPipeline(buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_graphicsPipeline);
}

void Pipeline::DefaultPipelineConfigInfo(PipelineConfigInfo& configInfo) {
    configInfo.name = "DefaultPipelineConfigInfo";

    configInfo.inputAssemblyInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    configInfo.inputAssemblyInfo.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    configInfo.inputAssemblyInfo.primitiveRestartEnable = VK_FALSE;

    configInfo.viewportInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    configInfo.viewportInfo.viewportCount = 1;
    configInfo.viewportInfo.pViewports = nullptr;
    configInfo.viewportInfo.scissorCount = 1;
    configInfo.viewportInfo.pScissors = nullptr;

    configInfo.rasterizationInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    configInfo.rasterizationInfo.depthClampEnable = VK_FALSE;
    configInfo.rasterizationInfo.rasterizerDiscardEnable = VK_FALSE;
    configInfo.rasterizationInfo.polygonMode = VK_POLYGON_MODE_FILL;
    configInfo.rasterizationInfo.lineWidth = 1.0f;
    configInfo.rasterizationInfo.cullMode = VK_CULL_MODE_BACK_BIT;
    configInfo.rasterizationInfo.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
    configInfo.rasterizationInfo.depthBiasEnable = VK_FALSE;
    configInfo.rasterizationInfo.depthBiasConstantFactor = 0.0f; // Optional
    configInfo.rasterizationInfo.depthBiasClamp = 0.0f; // Optional
    configInfo.rasterizationInfo.depthBiasSlopeFactor = 0.0f; // Optional

    configInfo.multisampleInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    configInfo.multisampleInfo.sampleShadingEnable = VK_FALSE;
    configInfo.multisampleInfo.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
    configInfo.multisampleInfo.minSampleShading = 1.0f; // Optional
    configInfo.multisampleInfo.pSampleMask = nullptr; // Optional
    configInfo.multisampleInfo.alphaToCoverageEnable = VK_FALSE; // Optional
    configInfo.multisampleInfo.alphaToOneEnable = VK_FALSE; // Optional

    configInfo.colorBlendAttachment.colorWriteMask =
            VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT |
            VK_COLOR_COMPONENT_A_BIT;
    configInfo.colorBlendAttachment.blendEnable = VK_FALSE;
    configInfo.colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_ONE; // Optional
    configInfo.colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ZERO; // Optional
    configInfo.colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD; // Optional
    configInfo.colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE; // Optional
    configInfo.colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO; // Optional
    configInfo.colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD; // Optional

    configInfo.colorBlendInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    configInfo.colorBlendInfo.logicOpEnable = VK_FALSE;
    configInfo.colorBlendInfo.logicOp = VK_LOGIC_OP_COPY; // Optional
    configInfo.colorBlendInfo.attachmentCount = 1;
    configInfo.colorBlendInfo.pAttachments = &configInfo.colorBlendAttachment;
    configInfo.colorBlendInfo.blendConstants[0] = 0.0f; // Optional
    configInfo.colorBlendInfo.blendConstants[1] = 0.0f; // Optional
    configInfo.colorBlendInfo.blendConstants[2] = 0.0f; // Optional
    configInfo.colorBlendInfo.blendConstants[3] = 0.0f; // Optional

    configInfo.depthStencilInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    configInfo.depthStencilInfo.depthTestEnable = VK_TRUE;
    configInfo.depthStencilInfo.depthWriteEnable = VK_TRUE;
    configInfo.depthStencilInfo.depthCompareOp = VK_COMPARE_OP_LESS;
    configInfo.depthStencilInfo.depthBoundsTestEnable = VK_FALSE;
    configInfo.depthStencilInfo.minDepthBounds = 0.0f; // Optional
    configInfo.depthStencilInfo.maxDepthBounds = 1.0f; // Optional
    configInfo.depthStencilInfo.stencilTestEnable = VK_FALSE;
    configInfo.depthStencilInfo.front = {}; // Optional
    configInfo.depthStencilInfo.back = {}; // Optional

    configInfo.dynamicStateEnables = {VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR, VK_DYNAMIC_STATE_POLYGON_MODE_EXT};
    configInfo.dynamicStateInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    configInfo.dynamicStateInfo.pDynamicStates = configInfo.dynamicStateEnables.data();
    configInfo.dynamicStateInfo.dynamicStateCount = static_cast<uint32_t>(configInfo.dynamicStateEnables.size());
    configInfo.dynamicStateInfo.flags = 0;

    configInfo.colorAttachments = {VK_FORMAT_B8G8R8A8_SRGB};
    configInfo.depthAttachment = VK_FORMAT_D32_SFLOAT;

    // configInfo.vertexAttributeDescriptions = std::move(Mesh::Vertex::getAttributeDescriptions());
    // configInfo.vertexBindingDescriptions = std::move(Mesh::Vertex::getBindingDescriptions());
}

std::vector<char> Pipeline::readFile(const std::string& filename) {
    std::ifstream file(filename, std::ios::ate | std::ios::binary);

    if (!file.is_open()) {
        throw std::runtime_error("Failed to open file: " + filename);
    }

    const size_t fileSize = (size_t)file.tellg();
    std::vector<char> buffer(fileSize);

    file.seekg(0);
    file.read(buffer.data(), static_cast<std::streamsize>(fileSize));

    file.close();

    return buffer;
}

void Pipeline::CreateGraphicsPipeline(const std::string& vertPath, const std::string& fragPath,
                                      const PipelineConfigInfo& configInfo) {
    assert(configInfo.pipelineLayout != VK_NULL_HANDLE && "no pipelineLayout provided in configInfo");

    std::vector < VkPipelineShaderStageCreateInfo > shaderStages;

    if (!vertPath.empty()) {
        std::string VertFileName = std::filesystem::path(vertPath).filename().string();
        auto vertCode = readFile(vertPath);
        spdlog::debug("Vertex shader code size: {}", vertCode.size());
        spdlog::debug("Vertex shader file: {}", VertFileName);

        CreateShaderModule(vertCode, &m_vertShaderModule);

        vkutil::addDebugLabel(
            m_device.getDevice(),
            m_vertShaderModule,
            VertFileName.c_str()
        );

        VkPipelineShaderStageCreateInfo vertShaderStageInfo{};
        vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
        vertShaderStageInfo.module = m_vertShaderModule;
        vertShaderStageInfo.pName = "main";
        vertShaderStageInfo.flags = 0;
        vertShaderStageInfo.pSpecializationInfo = nullptr;
        vertShaderStageInfo.pNext = nullptr;

        shaderStages.push_back(vertShaderStageInfo);
    }

    if (!fragPath.empty()) {
        std::string FragFileName = std::filesystem::path(fragPath).filename().string();
        auto fragCode = readFile(fragPath);

        spdlog::debug("Fragment shader code size: {}", fragCode.size());
        spdlog::debug("Fragment shader file: {}", FragFileName);

        CreateShaderModule(fragCode, &m_fragShaderModule);

        vkutil::addDebugLabel(
            m_device.getDevice(),
            m_fragShaderModule,
            FragFileName.c_str()
        );


        VkPipelineShaderStageCreateInfo fragShaderStageInfo{};
        fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
        fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
        fragShaderStageInfo.module = m_fragShaderModule;
        fragShaderStageInfo.pName = "main";
        fragShaderStageInfo.flags = 0;
        fragShaderStageInfo.pSpecializationInfo = nullptr;
        fragShaderStageInfo.pNext = nullptr;

        shaderStages.push_back(fragShaderStageInfo);
    }

    std::cout << std::endl;


    VkPipelineVertexInputStateCreateInfo vertexInputInfo{};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputInfo.vertexAttributeDescriptionCount = static_cast<uint32_t>(configInfo.vertexAttributeDescriptions.size());
    vertexInputInfo.pVertexAttributeDescriptions = configInfo.vertexAttributeDescriptions.data();
    vertexInputInfo.vertexBindingDescriptionCount = static_cast<uint32_t>(configInfo.vertexBindingDescriptions.size());
    vertexInputInfo.pVertexBindingDescriptions = configInfo.vertexBindingDescriptions.data();


    VkGraphicsPipelineCreateInfo pipelineInfo{};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = static_cast<uint32_t>(shaderStages.size());
    pipelineInfo.pStages = shaderStages.data();
    pipelineInfo.pVertexInputState = &vertexInputInfo;
    pipelineInfo.pInputAssemblyState = &configInfo.inputAssemblyInfo;
    pipelineInfo.pViewportState = &configInfo.viewportInfo;
    pipelineInfo.pRasterizationState = &configInfo.rasterizationInfo;
    pipelineInfo.pMultisampleState = &configInfo.multisampleInfo;
    pipelineInfo.pColorBlendState = &configInfo.colorBlendInfo;
    pipelineInfo.pDepthStencilState = &configInfo.depthStencilInfo;
    pipelineInfo.pDynamicState = &configInfo.dynamicStateInfo;

    VkPipelineRenderingCreateInfoKHR renderingInfo{};
    renderingInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO;
    renderingInfo.colorAttachmentCount = static_cast<uint32_t>(configInfo.colorAttachments.size());
    renderingInfo.pColorAttachmentFormats = configInfo.colorAttachments.data();
    renderingInfo.depthAttachmentFormat = configInfo.depthAttachment;

    pipelineInfo.pNext = &renderingInfo;

    pipelineInfo.layout = configInfo.pipelineLayout;
    pipelineInfo.renderPass = nullptr;
    // pipelineInfo.subpass = nullptr;

    pipelineInfo.basePipelineIndex = -1;
    pipelineInfo.basePipelineHandle = VK_NULL_HANDLE;


    if (vkCreateGraphicsPipelines(m_device.getDevice(), VK_NULL_HANDLE, 1, &pipelineInfo, nullptr, &m_graphicsPipeline) != VK_SUCCESS) {
        throw std::runtime_error("Can't make pipeline!");
    }

    if (!configInfo.name.empty()) {
        vkutil::addDebugLabel(m_device.getDevice(), m_graphicsPipeline, configInfo.name.c_str());

        const std::string layoutName = configInfo.name + " Layout";
        vkutil::addDebugLabel(m_device.getDevice(), configInfo.pipelineLayout, layoutName.c_str());
    }
}

void Pipeline::CreateShaderModule(const std::vector<char>& code, VkShaderModule* shaderModule) const {
    VkShaderModuleCreateInfo createInfo{};
    createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    createInfo.codeSize = code.size();
    createInfo.pCode = reinterpret_cast<const uint32_t*>(code.data());

    if (vkCreateShaderModule(m_device.getDevice(), &createInfo, nullptr, shaderModule) != VK_SUCCESS) {
        throw std::runtime_error("Failed to create shader module!");
    }
}