#include "imguipass.h"

#include <array>
#include <glm/glm.hpp>
#include <imgui.h>

#include "renderer.h"

ImGuiPass::ImGuiPass(Renderer& renderer) : renderer_(renderer) { 
    createDescriptorSetLayout();
    createPipeline();
    createFontImage();
}

ImGuiPass::~ImGuiPass() {
    vkDestroySampler(renderer_.getDevice(), fontSampler_, nullptr);
    vkDestroyImageView(renderer_.getDevice(), fontImageView_, nullptr);
    vkFreeMemory(renderer_.getDevice(), fontMemory_, nullptr);
    vkDestroyImage(renderer_.getDevice(), fontImage_, nullptr);
    
    vkDestroyPipeline(renderer_.getDevice(), pipeline_, nullptr);
    vkDestroyPipelineLayout(renderer_.getDevice(), pipelineLayout_, nullptr);
    
    vkDestroyDescriptorSetLayout(renderer_.getDevice(), setLayout_, nullptr);
}

void ImGuiPass::render(VkCommandBuffer commandBuffer, RenderTarget* target) {
    ImDrawData* drawData = ImGui::GetDrawData();
    
    VkBuffer& vertexBuffer = target->imguiVertexBuffers[target->currentResource];
    VkDeviceMemory& vertexMemory = target->imguiVertexMemorys[target->currentResource];
    
    VkBuffer& indexBuffer = target->imguiIndexBuffers[target->currentResource];
    VkDeviceMemory& indexMemory = target->imguiIndexMemorys[target->currentResource];
    
    const size_t vertexSize = drawData->TotalVtxCount * sizeof(ImDrawVert);
    if(vertexSize > target->imguiVertexBufferSizes[target->currentResource]) {
        createBuffer(vertexBuffer, vertexMemory, vertexSize, VK_BUFFER_USAGE_VERTEX_BUFFER_BIT);
        target->imguiVertexBufferSizes[target->currentResource] = vertexSize;
    }
    
    const size_t indexSize = drawData->TotalIdxCount * sizeof(ImDrawIdx);
    if(indexSize > target->imguiIndexBufferSizes[target->currentResource]) {
        createBuffer(indexBuffer, indexMemory, indexSize, VK_BUFFER_USAGE_INDEX_BUFFER_BIT);
        target->imguiIndexBufferSizes[target->currentResource] = indexSize;
    }
    
    if(vertexSize == 0 || indexSize == 0)
        return;
    
    ImDrawVert* vertexData = nullptr;
    ImDrawIdx* indexData = nullptr;
    vkMapMemory(renderer_.getDevice(), vertexMemory, 0, vertexSize, 0, reinterpret_cast<void**>(&vertexData));
    vkMapMemory(renderer_.getDevice(), indexMemory, 0, indexSize, 0, reinterpret_cast<void**>(&indexData));
    
    for(int i = 0; i < drawData->CmdListsCount; i++) {
        const ImDrawList* cmd_list = drawData->CmdLists[i];
        
        memcpy(vertexData, cmd_list->VtxBuffer.Data, cmd_list->VtxBuffer.Size * sizeof(ImDrawVert));
        vertexData += cmd_list->VtxBuffer.Size;
        
        memcpy(indexData, cmd_list->IdxBuffer.Data, cmd_list->IdxBuffer.Size * sizeof(ImDrawIdx));
        indexData += cmd_list->IdxBuffer.Size;
    }

    vkUnmapMemory(renderer_.getDevice(), vertexMemory);
    vkUnmapMemory(renderer_.getDevice(), indexMemory);
    
    vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_);
    
    VkDeviceSize offset = 0;
    vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offset);
    vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT16);
    
    float scale[2];
    scale[0] = 2.0f / drawData->DisplaySize.x;
    scale[1] = 2.0f / drawData->DisplaySize.y;
    
    float translate[2];
    translate[0] = -1.0f - drawData->DisplayPos.x * scale[0];
    translate[1] = -1.0f - drawData->DisplayPos.y * scale[1];
    
    vkCmdPushConstants(commandBuffer, pipelineLayout_, VK_SHADER_STAGE_VERTEX_BIT, sizeof(float) * 0, sizeof(float) * 2, scale);
    vkCmdPushConstants(commandBuffer, pipelineLayout_, VK_SHADER_STAGE_VERTEX_BIT, sizeof(float) * 2, sizeof(float) * 2, translate);
    
    int vertexOffset = 0, indexOffset = 0;
    const ImVec2 displayPos = drawData->DisplayPos;
    for(int n = 0; n < drawData->CmdListsCount; n++) {
        const ImDrawList* cmd_list = drawData->CmdLists[n];
        for(int cmd_i = 0; cmd_i < cmd_list->CmdBuffer.Size; cmd_i++) {
            const ImDrawCmd* pcmd = &cmd_list->CmdBuffer[cmd_i];
            
            if(descriptorSets_.count((VkImageView)pcmd->TextureId)) {
                vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout_, 0, 1, &descriptorSets_[(VkImageView)pcmd->TextureId], 0, nullptr);
            } else {
                VkDescriptorSetAllocateInfo allocInfo = {};
                allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
                allocInfo.descriptorPool = renderer_.getDescriptorPool();
                allocInfo.descriptorSetCount = 1;
                allocInfo.pSetLayouts = &setLayout_;
                
                VkDescriptorSet set = nullptr;
                vkAllocateDescriptorSets(renderer_.getDevice(), &allocInfo, &set);
                
                VkDescriptorImageInfo imageInfo = {};
                imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
                imageInfo.imageView = (VkImageView)pcmd->TextureId;
                imageInfo.sampler = fontSampler_;
                
                VkWriteDescriptorSet descriptorWrite = {};
                descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
                descriptorWrite.descriptorCount = 1;
                descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
                descriptorWrite.dstSet = set;
                descriptorWrite.pImageInfo = &imageInfo;
                
                vkUpdateDescriptorSets(renderer_.getDevice(), 1, &descriptorWrite, 0, nullptr);
                
                vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout_, 0, 1, &set, 0, nullptr);
                
                descriptorSets_[(VkImageView)pcmd->TextureId] = set;
            }
                
            if(pcmd->UserCallback) {
                pcmd->UserCallback(cmd_list, pcmd);
            } else {
                VkRect2D scissor;
                scissor.offset.x = (int32_t)(pcmd->ClipRect.x - displayPos.x) > 0 ? (int32_t)(pcmd->ClipRect.x - displayPos.x) : 0;
                scissor.offset.y = (int32_t)(pcmd->ClipRect.y - displayPos.y) > 0 ? (int32_t)(pcmd->ClipRect.y - displayPos.y) : 0;
                scissor.extent.width = (uint32_t)(pcmd->ClipRect.z - pcmd->ClipRect.x);
                scissor.extent.height = (uint32_t)(pcmd->ClipRect.w - pcmd->ClipRect.y + 1);
                
                vkCmdSetScissor(commandBuffer, 0, 1, &scissor);
                
                vkCmdDrawIndexed(commandBuffer, pcmd->ElemCount, 1, indexOffset, vertexOffset, 0);
            }
            
            indexOffset += pcmd->ElemCount;
        }
        
        vertexOffset += cmd_list->VtxBuffer.Size;
    }
}

void ImGuiPass::createDescriptorSetLayout() {
    VkDescriptorSetLayoutBinding samplerBinding = {};
    samplerBinding.descriptorCount = 1;
    samplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    samplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;
    
    VkDescriptorSetLayoutCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    createInfo.bindingCount = 1;
    createInfo.pBindings = &samplerBinding;
    
    vkCreateDescriptorSetLayout(renderer_.getDevice(), &createInfo, nullptr, &setLayout_);
}

void ImGuiPass::createPipeline() {
    VkShaderModule vertShaderModule = renderer_.createShader("shaders/imgui.vert.spv");
    VkShaderModule fragShaderModule = renderer_.createShader("shaders/imgui.frag.spv");
    
    VkPipelineShaderStageCreateInfo vertShaderStageInfo = {};
    vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
    vertShaderStageInfo.module = vertShaderModule;
    vertShaderStageInfo.pName = "main";
    
    VkPipelineShaderStageCreateInfo fragShaderStageInfo = {};
    fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    fragShaderStageInfo.module = fragShaderModule;
    fragShaderStageInfo.pName = "main";
    
    const std::array<VkPipelineShaderStageCreateInfo, 2> shaderStages = {vertShaderStageInfo, fragShaderStageInfo};
    
    VkVertexInputBindingDescription vertexBindingDescription = {};
    vertexBindingDescription.stride = sizeof(ImDrawVert);
    vertexBindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
    
    VkVertexInputAttributeDescription positionAttribute = {};
    positionAttribute.format = VK_FORMAT_R32G32_SFLOAT;
    positionAttribute.offset = offsetof(ImDrawVert, pos);
    
    VkVertexInputAttributeDescription uvAttribute = {};
    uvAttribute.location = 1;
    uvAttribute.format = VK_FORMAT_R32G32_SFLOAT;
    uvAttribute.offset = offsetof(ImDrawVert, uv);
    
    VkVertexInputAttributeDescription colorAttribute = {};
    colorAttribute.location = 2;
    colorAttribute.format = VK_FORMAT_R8G8B8A8_UNORM;
    colorAttribute.offset = offsetof(ImDrawVert, col);
    
    const std::array<VkVertexInputAttributeDescription, 3> attributes = {
        positionAttribute,
        uvAttribute,
        colorAttribute
    };

    VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputInfo.vertexBindingDescriptionCount = 1;
    vertexInputInfo.pVertexBindingDescriptions = &vertexBindingDescription;
    vertexInputInfo.vertexAttributeDescriptionCount = attributes.size();
    vertexInputInfo.pVertexAttributeDescriptions = attributes.data();
    
    VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
    inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
    
    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.polygonMode = VK_POLYGON_MODE_FILL;
    rasterizer.cullMode = VK_CULL_MODE_NONE;
    rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
    rasterizer.lineWidth = 1.0f;
    
    VkPipelineMultisampleStateCreateInfo multisampling = {};
    multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisampling.rasterizationSamples = VK_SAMPLE_COUNT_1_BIT;
    
    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_MINUS_SRC_ALPHA;
    colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
    colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;
    
    VkPipelineColorBlendStateCreateInfo colorBlending = {};
    colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlending.attachmentCount = 1;
    colorBlending.pAttachments = &colorBlendAttachment;
    
    const std::array<VkDynamicState, 2> dynamicStates = {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR
    };
    
    VkPipelineDynamicStateCreateInfo dynamicState = {};
    dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicState.dynamicStateCount = dynamicStates.size();
    dynamicState.pDynamicStates = dynamicStates.data();
    
    VkPushConstantRange pushConstant = {};
    pushConstant.size = sizeof(glm::vec4);
    pushConstant.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;
    
    VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    pipelineLayoutInfo.setLayoutCount = 1;
    pipelineLayoutInfo.pSetLayouts = &setLayout_;
    pipelineLayoutInfo.pushConstantRangeCount = 1;
    pipelineLayoutInfo.pPushConstantRanges = &pushConstant;
    
    vkCreatePipelineLayout(renderer_.getDevice(), &pipelineLayoutInfo, nullptr, &pipelineLayout_);
    
    VkGraphicsPipelineCreateInfo pipelineInfo = {};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = shaderStages.size();
    pipelineInfo.pStages = shaderStages.data();
    pipelineInfo.pVertexInputState = &vertexInputInfo;
    pipelineInfo.pInputAssemblyState = &inputAssembly;
    pipelineInfo.pViewportState = &viewportState;
    pipelineInfo.pRasterizationState = &rasterizer;
    pipelineInfo.pMultisampleState = &multisampling;
    pipelineInfo.pColorBlendState = &colorBlending;
    pipelineInfo.pDynamicState = &dynamicState;
    pipelineInfo.layout = pipelineLayout_;
    pipelineInfo.renderPass = renderer_.getRenderPass();
    
    vkCreateGraphicsPipelines(renderer_.getDevice(), nullptr, 1, &pipelineInfo, nullptr, &pipeline_);
    
    vkDestroyShaderModule(renderer_.getDevice(), fragShaderModule, nullptr);
    vkDestroyShaderModule(renderer_.getDevice(), vertShaderModule, nullptr);
}

void ImGuiPass::createFontImage() {
    ImGuiIO& io = ImGui::GetIO();
    
    unsigned char* pixels = nullptr;
    int width = 0, height = 0;
    io.Fonts->GetTexDataAsRGBA32(&pixels, &width, &height);
    
    VkImageCreateInfo imageCreateInfo = {};
    imageCreateInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageCreateInfo.imageType = VK_IMAGE_TYPE_2D;
    imageCreateInfo.extent.width = width;
    imageCreateInfo.extent.height = height;
    imageCreateInfo.extent.depth = 1;
    imageCreateInfo.mipLevels = 1;
    imageCreateInfo.arrayLayers = 1;
    imageCreateInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    imageCreateInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageCreateInfo.usage = VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT;
    imageCreateInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    imageCreateInfo.samples = VK_SAMPLE_COUNT_1_BIT;
    
    vkCreateImage(renderer_.getDevice(), &imageCreateInfo, nullptr, &fontImage_);
    
    VkMemoryRequirements memRequirements;
    vkGetImageMemoryRequirements(renderer_.getDevice(), fontImage_, &memRequirements);
    
    VkMemoryAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = renderer_.findMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    
    vkAllocateMemory(renderer_.getDevice(), &allocInfo, nullptr, &fontMemory_);
    vkBindImageMemory(renderer_.getDevice(), fontImage_, fontMemory_, 0);
 
    renderer_.uploadImageData(
        fontImage_,
        width,
        height,
        width * height * 4,
        pixels
    );
        
    VkImageViewCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    createInfo.image = fontImage_;
    createInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    createInfo.format = VK_FORMAT_R8G8B8A8_UNORM;
    createInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    createInfo.subresourceRange.levelCount = 1;
    createInfo.subresourceRange.layerCount = 1;
    
    vkCreateImageView(renderer_.getDevice(), &createInfo, nullptr, &fontImageView_);
    
    io.Fonts->TexID = static_cast<ImTextureID>(fontImageView_);
    
    VkSamplerCreateInfo samplerInfo = {};
    samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
    samplerInfo.magFilter = VK_FILTER_LINEAR;
    samplerInfo.minFilter = VK_FILTER_LINEAR;
    samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_BORDER;
    samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
    
    vkCreateSampler(renderer_.getDevice(), &samplerInfo, nullptr, &fontSampler_);
}

void ImGuiPass::createBuffer(VkBuffer& buffer, VkDeviceMemory& memory, VkDeviceSize size, VkBufferUsageFlagBits bufferUsage) {
    if(buffer != nullptr)
        vkDestroyBuffer(renderer_.getDevice(), buffer, nullptr);
    
    if(memory != nullptr)
        vkFreeMemory(renderer_.getDevice(), memory, nullptr);
    
    VkBufferCreateInfo bufferInfo = {};
    bufferInfo.sType = VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO;
    bufferInfo.size = size;
    bufferInfo.usage = bufferUsage;
    bufferInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    
    vkCreateBuffer(renderer_.getDevice(), &bufferInfo, nullptr, &buffer);
    
    VkMemoryRequirements memRequirements = {};
    vkGetBufferMemoryRequirements(renderer_.getDevice(), buffer, &memRequirements);
    
    VkMemoryAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memRequirements.size;
    allocInfo.memoryTypeIndex = renderer_.findMemoryType(memRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT);
    
    vkAllocateMemory(renderer_.getDevice(), &allocInfo, nullptr, &memory);
    vkBindBufferMemory(renderer_.getDevice(), buffer, memory, 0);
}