#include "postpass.h"

#include <array>

#include "renderer.h"

PostPass::PostPass(Renderer& renderer) : renderer_(renderer) {
    createDescriptorSetLayout();
    createPipeline();
    createSampler();
}

PostPass::~PostPass() {
    vkDestroySampler(renderer_.getDevice(), offscreenSampler_, nullptr);

    vkDestroyPipelineLayout(renderer_.getDevice(), pipelineLayout_, nullptr);
    vkDestroyPipeline(renderer_.getDevice(), pipeline_, nullptr);

    vkDestroyDescriptorSetLayout(renderer_.getDevice(), setLayout_, nullptr);
}

void PostPass::render(VkCommandBuffer commandBuffer, RenderTarget* target) {
    vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipeline_);
    vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, pipelineLayout_, 0, 1, &target->postSets[target->currentResource], 0, nullptr);

    glm::vec4 viewport;
    viewport.x = 1.0 / target->extent.width;
    viewport.y = 1.0 / target->extent.height;
    viewport.z = target->extent.width;
    viewport.w = target->extent.height;

    vkCmdPushConstants(commandBuffer, pipelineLayout_, VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(glm::vec4), &viewport);

    vkCmdDraw(commandBuffer, 3, 1, 0, 0);
}

void PostPass::createDescriptorSet(RenderTarget* target) {
    VkDescriptorSetAllocateInfo allocInfo = {};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = renderer_.getDescriptorPool();
    allocInfo.descriptorSetCount = target->numImages;

    // FIXME: lol what
    auto layouts = new VkDescriptorSetLayout[target->numImages];
    for(uint32_t i = 0; i < target->numImages; i++)
        layouts[i] = setLayout_;

    allocInfo.pSetLayouts = layouts;

    target->postSets = new VkDescriptorSet[target->numImages];
    vkAllocateDescriptorSets(renderer_.getDevice(), &allocInfo, target->postSets);

    delete[] layouts;

    for(uint32_t i = 0; i < target->numImages; i++) {
        VkDescriptorImageInfo sceneImageInfo = {};
        sceneImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        sceneImageInfo.imageView = target->offscreenColorImageViews[i];
        sceneImageInfo.sampler = offscreenSampler_;

        VkDescriptorImageInfo depthImageInfo = {};
        depthImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;;
        depthImageInfo.imageView = target->offscreenDepthImageViews[i];
        depthImageInfo.sampler = offscreenSampler_;

        VkDescriptorImageInfo nearFieldImageInfo = {};
        nearFieldImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        nearFieldImageInfo.imageView = target->nearFieldImageViews[i];
        nearFieldImageInfo.sampler = offscreenSampler_;

        VkDescriptorImageInfo farFieldImageInfo = {};
        farFieldImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        farFieldImageInfo.imageView = target->farFieldImageViews[i];
        farFieldImageInfo.sampler = offscreenSampler_;

        VkDescriptorImageInfo blendImageInfo = {};
        blendImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        blendImageInfo.imageView = target->blendImageViews[i];
        blendImageInfo.sampler = offscreenSampler_;

        VkDescriptorImageInfo sobelImageInfo = {};
        sobelImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        sobelImageInfo.imageView = target->sobelImageViews[i];
        sobelImageInfo.sampler = offscreenSampler_;

        VkWriteDescriptorSet sceneDescriptorWrite = {};
        sceneDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        sceneDescriptorWrite.descriptorCount = 1;
        sceneDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        sceneDescriptorWrite.dstSet = target->postSets[i];
        sceneDescriptorWrite.pImageInfo = &sceneImageInfo;

        VkWriteDescriptorSet depthDescriptorWrite = {};
        depthDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        depthDescriptorWrite.descriptorCount = 1;
        depthDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        depthDescriptorWrite.dstBinding = 1;
        depthDescriptorWrite.dstSet = target->postSets[i];
        depthDescriptorWrite.pImageInfo = &depthImageInfo;

        VkWriteDescriptorSet nearFieldDescriptorWrite = {};
        nearFieldDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        nearFieldDescriptorWrite.descriptorCount = 1;
        nearFieldDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        nearFieldDescriptorWrite.dstBinding = 2;
        nearFieldDescriptorWrite.dstSet = target->postSets[i];
        nearFieldDescriptorWrite.pImageInfo = &nearFieldImageInfo;

        VkWriteDescriptorSet farFieldDescriptorWrite = {};
        farFieldDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        farFieldDescriptorWrite.descriptorCount = 1;
        farFieldDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        farFieldDescriptorWrite.dstBinding = 3;
        farFieldDescriptorWrite.dstSet = target->postSets[i];
        farFieldDescriptorWrite.pImageInfo = &farFieldImageInfo;

        VkWriteDescriptorSet blendDescriptorWrite = {};
        blendDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        blendDescriptorWrite.descriptorCount = 1;
        blendDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        blendDescriptorWrite.dstBinding = 4;
        blendDescriptorWrite.dstSet = target->postSets[i];
        blendDescriptorWrite.pImageInfo = &blendImageInfo;

        VkWriteDescriptorSet sobelDescriptorWrite = {};
        sobelDescriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
        sobelDescriptorWrite.descriptorCount = 1;
        sobelDescriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        sobelDescriptorWrite.dstBinding = 5;
        sobelDescriptorWrite.dstSet = target->postSets[i];
        sobelDescriptorWrite.pImageInfo = &sobelImageInfo;

        const std::array<VkWriteDescriptorSet, 6> descriptorWrites = {
            sceneDescriptorWrite,
            depthDescriptorWrite,
            nearFieldDescriptorWrite,
            farFieldDescriptorWrite,
            blendDescriptorWrite,
            sobelDescriptorWrite
        };

        vkUpdateDescriptorSets(renderer_.getDevice(), descriptorWrites.size(), descriptorWrites.data(), 0, nullptr);
    }
}

void PostPass::createDescriptorSetLayout() {
    VkDescriptorSetLayoutBinding offscreenSamplerBinding = {};
    offscreenSamplerBinding.descriptorCount = 1;
    offscreenSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    offscreenSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding depthSamplerBinding = {};
    depthSamplerBinding.binding = 1;
    depthSamplerBinding.descriptorCount = 1;
    depthSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    depthSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding nearFieldSamplerBinding = {};
    nearFieldSamplerBinding.binding = 2;
    nearFieldSamplerBinding.descriptorCount = 1;
    nearFieldSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    nearFieldSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding farFieldSamplerBinding = {};
    farFieldSamplerBinding.binding = 3;
    farFieldSamplerBinding.descriptorCount = 1;
    farFieldSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    farFieldSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding blendSamplerBinding = {};
    blendSamplerBinding.binding = 4;
    blendSamplerBinding.descriptorCount = 1;
    blendSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    blendSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    VkDescriptorSetLayoutBinding sobelSamplerBinding = {};
    sobelSamplerBinding.binding = 5;
    sobelSamplerBinding.descriptorCount = 1;
    sobelSamplerBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    sobelSamplerBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

    const std::array<VkDescriptorSetLayoutBinding, 6> bindings = {
        offscreenSamplerBinding,
        depthSamplerBinding,
        nearFieldSamplerBinding,
        farFieldSamplerBinding,
        blendSamplerBinding,
        sobelSamplerBinding
    };

    VkDescriptorSetLayoutCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
    createInfo.bindingCount = bindings.size();
    createInfo.pBindings = bindings.data();

    vkCreateDescriptorSetLayout(renderer_.getDevice(), &createInfo, nullptr, &setLayout_);
}

void PostPass::createPipeline() {
    VkShaderModule vertShaderModule = renderer_.createShader("shaders/post.vert.spv");
    VkShaderModule fragShaderModule = renderer_.createShader("shaders/post.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};

    VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;

    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.lineWidth = 1.0f;
    rasterizer.cullMode = VK_CULL_MODE_FRONT_BIT;
    rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;

    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_FALSE;

    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_FRAGMENT_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 PostPass::createSampler() {
    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_REPEAT;
    samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
    samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
    samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
    samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;

    vkCreateSampler(renderer_.getDevice(), &samplerInfo, nullptr, &offscreenSampler_);
}