#include "CommandCenter.hpp"

#include "PhysicalDeviceInfo.hpp"
#include "Pipeline.hpp"
#include "Renderer.hpp"
#include <TracyVulkan.hpp>

using namespace render::vk;

#define CONTENT_DIR "content/"
#define TEXTURES_DIR CONTENT_DIR "textures/"

CommandCenter::CommandCenter(VkDevice device, const PhysicalDeviceInfo &info, const render::renderOptions &opt): device(device) {
    { // Graphics command pool
        vkGetDeviceQueue(device, info.queueIndices.graphicsFamily.value(), 0, &graphicsQueue);
        VkCommandPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
        poolInfo.queueFamilyIndex = info.queueIndices.graphicsFamily.value();
        poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;

        if (vkCreateCommandPool(device, &poolInfo, ALLOC, &graphicsPool) != VK_SUCCESS) {
            FATAL("Failed to create graphics pool!");
        }
    }

    indicSphereBuffer = Shape::Create(Shape::LINE_SPHERE);
    indicCubeBuffer = Shape::Create(Shape::LINE_CUBE);
    if (!(indicCubeBuffer && indicSphereBuffer)) {
        FATAL("Failed to create vertex buffer!");
    }

    skyCubeBuffer = Shape::Create(Shape::SKY_CUBE);
    if (!skyCubeBuffer) {
        FATAL("Failed to create vertex buffer!");
    }
    skyboxTexture = TextureCube::LoadFromFiles(TEXTURES_DIR + opt.textures + "/sky/Space_tray.cube", {});
    if (!skyboxTexture) {
        FATAL("Failed to create texture sampler!");
    }

    colorFormat = info.getSurfaceFormat().format;
    colorSamples = info.samples;
    LOG_D("Samples: " << colorSamples);
    depthFormat = info.findDepthFormat();

    loadAtlases(opt.textures, opt.getAnisotropy(), opt.getMipmapLodBias());
}
CommandCenter::~CommandCenter() {
    if(!freed)
        free();

    vkDestroyCommandPool(device, graphicsPool, ALLOC);
}
#include "../../../core/world/materials.hpp"
void CommandCenter::loadAtlases(const std::string& textures, int anisotropy, float lodBias) {
    voxelTextureAtlas.reset();
    voxelNormalAtlas.reset();
    voxelHOSAtlas.reset();

    std::vector<std::string> paths;
    auto makePaths = [&](const std::string &suffix) {
        paths.clear();
        paths.reserve(world::materials::textures.size());
        for (auto& texture: world::materials::textures) {
            paths.push_back(TEXTURES_DIR + textures + "/terrain/" + texture + suffix + ".dds");
        }
        return paths;
    };
    const auto sampling = Texture::sampling{true, true, Texture::Wrap::REPEAT, anisotropy, true, lodBias};
    voxelTextureAtlas = TextureArray::LoadFromFiles(makePaths(""), sampling);
    voxelNormalAtlas = TextureArray::LoadFromFiles(makePaths(".nrm"), sampling, false);
    voxelHOSAtlas = TextureArray::LoadFromFiles(makePaths(".hos"), sampling);

    if(!(voxelTextureAtlas && voxelNormalAtlas && voxelHOSAtlas)) {
        FATAL("Failed to load texture pack!");
    }
}

#include <glm/gtc/matrix_transform.hpp>
void CommandCenter::allocate(const std::vector<VkImageView>& views, const Pipeline& pipe, VkPhysicalDevice physicalDevice, VkExtent2D extent) {
    assert(freed);

    if (colorSamples > 1) {
        colorbuffer = Image::Create(Image::requirement({{extent.height, extent.width}, 1, (Image::Format)colorFormat},
            Image::Layout::COLOR_ATTACHMENT, Image::Usage::COLOR_ATTACHMENT | Image::Usage::TRANSIENT_ATTACHMENT, // NOTE: VulkanTutorial#118
            Image::Aspect::COLOR, colorSamples));
    }

    depthbuffer = Image::Create(Image::requirement({{extent.height, extent.width}, 1, (Image::Format)depthFormat},
        Image::Layout::DEPTH_STENCIL_ATTACHMENT, Image::Usage::DEPTH_STENCIL_ATTACHMENT,
        Image::Aspect::DEPTH, colorSamples));

    framebuffers.resize(views.size());

    for (size_t i = 0; i < views.size(); i++) {
        std::vector<VkImageView> attachments;
        if (colorSamples > 1) {
            attachments = {colorbuffer->getView(), depthbuffer->getView(), views[i]};
        } else {
            attachments = {views[i], depthbuffer->getView()};
        }

        VkFramebufferCreateInfo framebufferInfo{};
        framebufferInfo.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
        framebufferInfo.renderPass = pipe.getRenderPass();
        framebufferInfo.attachmentCount = attachments.size();
        framebufferInfo.pAttachments = attachments.data();
        framebufferInfo.width = extent.width;
        framebufferInfo.height = extent.height;
        framebufferInfo.layers = 1;

        if (vkCreateFramebuffer(device, &framebufferInfo, ALLOC, &framebuffers[i]) != VK_SUCCESS) {
            FATAL("Failed to create framebuffer!");
        }
    }

    uiFramebuffers.resize(views.size());

    for (size_t i = 0; i < views.size(); i++) {
        VkFramebufferCreateInfo info = {};
        info.sType = VK_STRUCTURE_TYPE_FRAMEBUFFER_CREATE_INFO;
        info.renderPass = pipe.getUIRenderPass();
        info.attachmentCount = 1;
        info.pAttachments = &views[i];
        info.width = extent.width;
        info.height = extent.height;
        info.layers = 1;

        if (vkCreateFramebuffer(device, &info, ALLOC, &uiFramebuffers[i]) != VK_SUCCESS) {
            FATAL("Failed to create ui framebuffer!");
        }
    }

    { // Uniform buffers
        std::vector<Buffer::requirement> requirements;
        requirements.resize(framebuffers.size(), Buffer::requirement(sizeof(VoxelUBO), Buffer::Usage::UNIFORM));
        uniformBuffers.allocate(requirements, true);
        if (!uniformBuffers) {
            FATAL("Failed to allocate UBO");
        }
    }
    { // Descriptor pools
        std::array<VkDescriptorPoolSize, 2> poolSizes{};
        poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[0].descriptorCount = framebuffers.size() * 3;
        poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        poolSizes[1].descriptorCount = framebuffers.size() * 4;

        VkDescriptorPoolCreateInfo poolInfo{};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.poolSizeCount = poolSizes.size();
        poolInfo.pPoolSizes = poolSizes.data();
        poolInfo.maxSets = framebuffers.size() * 3;
        poolInfo.flags = 0; //VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT

        if (vkCreateDescriptorPool(device, &poolInfo, nullptr, &descriptorPool) != VK_SUCCESS) {
            FATAL("Failed to create descriptor pool!");
        }
    }
    { // Descriptor sets
        VkDescriptorSetAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        allocInfo.descriptorSetCount = framebuffers.size();

        auto allocSets = [&](VkDescriptorSetLayout layout, std::vector<VkDescriptorSet>& out) {
            std::vector<VkDescriptorSetLayout> layouts(framebuffers.size(), layout);
            allocInfo.descriptorPool = descriptorPool;
            allocInfo.pSetLayouts = layouts.data();

            out.resize(framebuffers.size());
            if (vkAllocateDescriptorSets(device, &allocInfo, out.data()) != VK_SUCCESS) {
                FATAL("Failed to allocate descriptor sets!");
            }
        };
        allocSets(pipe.getVoxelDescriptorSet(), voxelDescriptorSets);
        allocSets(pipe.getIndicDescriptorSet(), indicDescriptorSets);
        allocSets(pipe.getSkyDescriptorSet(), skyDescriptorSets);

        for (size_t i = 0; i < voxelDescriptorSets.size(); i++) {
            VkDescriptorBufferInfo bufferInfo{};
            bufferInfo.buffer = uniformBuffers.at(i);
            bufferInfo.offset = 0;
            bufferInfo.range = sizeof(VoxelUBO);

            std::array<VkWriteDescriptorSet, 4> descriptorWrites{};
            descriptorWrites[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[0].dstSet = voxelDescriptorSets[i];
            descriptorWrites[0].dstBinding = 0;
            descriptorWrites[0].dstArrayElement = 0;
            descriptorWrites[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
            descriptorWrites[0].descriptorCount = 1;
            descriptorWrites[0].pBufferInfo = &bufferInfo;

            descriptorWrites[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[1].dstSet = voxelDescriptorSets[i];
            descriptorWrites[1].dstBinding = 1;
            descriptorWrites[1].dstArrayElement = 0;
            descriptorWrites[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            descriptorWrites[1].descriptorCount = 1;
            descriptorWrites[1].pImageInfo = &voxelTextureAtlas->getDescriptor();
            descriptorWrites[2].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[2].dstSet = voxelDescriptorSets[i];
            descriptorWrites[2].dstBinding = 2;
            descriptorWrites[2].dstArrayElement = 0;
            descriptorWrites[2].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            descriptorWrites[2].descriptorCount = 1;
            descriptorWrites[2].pImageInfo = &voxelNormalAtlas->getDescriptor();
            descriptorWrites[3].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrites[3].dstSet = voxelDescriptorSets[i];
            descriptorWrites[3].dstBinding = 3;
            descriptorWrites[3].dstArrayElement = 0;
            descriptorWrites[3].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            descriptorWrites[3].descriptorCount = 1;
            descriptorWrites[3].pImageInfo = &voxelHOSAtlas->getDescriptor();

            vkUpdateDescriptorSets(device, descriptorWrites.size(), descriptorWrites.data(), 0, nullptr);

            bufferInfo.range = sizeof(ViewProjUBO);

            descriptorWrites[0].dstSet = skyDescriptorSets[i];
            descriptorWrites[1].dstSet = skyDescriptorSets[i];
            descriptorWrites[1].pImageInfo = &skyboxTexture->getDescriptor();
            vkUpdateDescriptorSets(device, 2, descriptorWrites.data(), 0, nullptr);

            descriptorWrites[0].dstSet = indicDescriptorSets[i];
            vkUpdateDescriptorSets(device, 1, descriptorWrites.data(), 0, nullptr);
        }
    }
    {
        graphicsBuffers.resize(framebuffers.size());
        VkCommandBufferAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        allocInfo.commandPool = graphicsPool;
        allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        allocInfo.commandBufferCount = (uint32_t) graphicsBuffers.size();

        if (vkAllocateCommandBuffers(device, &allocInfo, graphicsBuffers.data()) != VK_SUCCESS) {
            FATAL("Failed to allocate graphics buffers!");
        }

        (void)physicalDevice;
        tracyCtx = TracyVkContext(physicalDevice, device, graphicsQueue, graphicsBuffers.front());
    }

    freed = false;
}
void CommandCenter::free() {
    assert(!freed);
    TracyVkDestroy(tracyCtx);
    vkFreeCommandBuffers(device, graphicsPool, static_cast<uint32_t>(graphicsBuffers.size()), graphicsBuffers.data());

    vkDestroyDescriptorPool(device, descriptorPool, ALLOC);

    colorbuffer.reset();
    depthbuffer.reset();

    for (size_t i = 0; i < framebuffers.size(); i++) {
        vkDestroyFramebuffer(device, framebuffers[i], ALLOC);
    }
    for (size_t i = 0; i < uiFramebuffers.size(); i++) {
        vkDestroyFramebuffer(device, uiFramebuffers[i], ALLOC);
    }

    freed = true;
}

#include <chrono>
#include <memory.h>

void CommandCenter::startRecording(uint32_t idx, VkRenderPass renderPass, VkExtent2D extent, const VoxelUBO& ubo) {
    uniformBuffers.write(idx, data_view(&ubo, sizeof(ubo)));


    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
    //TODO: reuse
    beginInfo.pInheritanceInfo = nullptr;

    if (vkBeginCommandBuffer(graphicsBuffers[idx], &beginInfo) != VK_SUCCESS) {
        FATAL("Failed to begin recording command buffer!");
    }

    if (!idx) {
        TracyVkCollect(tracyCtx, graphicsBuffers[idx]);
    }
    TracyVkZone(tracyCtx, graphicsBuffers[idx], "Begin");
    VkRenderPassBeginInfo renderPassInfo{};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    renderPassInfo.renderPass = renderPass;
    renderPassInfo.framebuffer = framebuffers[idx];
    renderPassInfo.renderArea.offset = {0, 0};
    renderPassInfo.renderArea.extent = extent;

    std::array<VkClearValue, 2> clearValues{};
    clearValues[0].color = {ubo.fog.x, ubo.fog.y, ubo.fog.z, 1};
    clearValues[1].depthStencil = {1.0f, 0};
    renderPassInfo.clearValueCount = clearValues.size();
    renderPassInfo.pClearValues = clearValues.data();

    vkCmdBeginRenderPass(graphicsBuffers[idx], &renderPassInfo, VK_SUBPASS_CONTENTS_INLINE);
}

void CommandCenter::startWorldPass(uint32_t idx, const Subpass &worldPass) {
    vkCmdBindPipeline(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, worldPass.pipeline);
    vkCmdBindDescriptorSets(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, worldPass.layout, 0, 1, &voxelDescriptorSets[idx], 0, nullptr);
}
size_t CommandCenter::recordModel(uint32_t i, const Subpass &worldPass, const UniqueCurvaturePush& push, const LodModel *const modelBuffer) {
    vkCmdPushConstants(graphicsBuffers[i], worldPass.layout, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(push), &push);

    VkBuffer vertexBuffers[] = {modelBuffer->getVertex()};
    VkDeviceSize offsets[] = {0};
    vkCmdBindVertexBuffers(graphicsBuffers[i], 0, 1, vertexBuffers, offsets);
    vkCmdBindIndexBuffer(graphicsBuffers[i], modelBuffer->getIndex(), 0, VK_INDEX_TYPE_UINT16);
    auto size = modelBuffer->getIndexSize();
    vkCmdDrawIndexed(graphicsBuffers[i], size, 1, modelBuffer->getIndexStart(), 0, 0);
    return size;
}
void CommandCenter::startEntityPass(uint32_t, const Subpass&) { }
size_t CommandCenter::recordModels(uint32_t i, const Subpass &entityPass, const std::vector<glm::mat4> &matrices, const Model *const modelBuffer) {
    VkBuffer vertexBuffers[] = {modelBuffer->getVertex()};
    VkDeviceSize offsets[] = {0};
    vkCmdBindVertexBuffers(graphicsBuffers[i], 0, 1, vertexBuffers, offsets);
    vkCmdBindIndexBuffer(graphicsBuffers[i], modelBuffer->getIndex(), 0, VK_INDEX_TYPE_UINT16);

    auto size = modelBuffer->getIndexSize();
    UniqueCurvaturePush push;
    push.curvature = 0;
    for(const auto& matrix: matrices) {
        push.model = matrix;
        vkCmdPushConstants(graphicsBuffers[i], entityPass.layout, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(push), &push);
        vkCmdDrawIndexed(graphicsBuffers[i], size, 1, 0, 0, 0);
    }
    return size * matrices.size();
}
void CommandCenter::startIndicPass(uint32_t idx, const Subpass& indicPass) {
    vkCmdNextSubpass(graphicsBuffers[idx], VK_SUBPASS_CONTENTS_INLINE);
    vkCmdBindPipeline(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, indicPass.pipeline);
    vkCmdBindDescriptorSets(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, indicPass.layout, 0, 1, &indicDescriptorSets[idx], 0, nullptr);
}
size_t CommandCenter::recordIndicator(uint32_t idx, const Subpass& indicPass, const ModelColorPush& push, bool isCube) {
    vkCmdPushConstants(graphicsBuffers[idx], indicPass.layout, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(push), &push);

    const auto buffer = isCube ? indicCubeBuffer.get() : indicSphereBuffer.get();
    VkBuffer vertexBuffers[] = {buffer->getRef()};
    VkDeviceSize offsets[] = {0};
    vkCmdBindVertexBuffers(graphicsBuffers[idx], 0, 1, vertexBuffers, offsets);
    vkCmdDraw(graphicsBuffers[idx], buffer->size, 1, 0, 0);
    return buffer->size;
}
void CommandCenter::recordPostprocess(uint32_t idx, const Subpass& skyPass, bool skybox) {
    vkCmdNextSubpass(graphicsBuffers[idx], VK_SUBPASS_CONTENTS_INLINE);
    if (skybox) {
        vkCmdBindPipeline(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, skyPass.pipeline);
        vkCmdBindDescriptorSets(graphicsBuffers[idx], VK_PIPELINE_BIND_POINT_GRAPHICS, skyPass.layout, 0, 1, &skyDescriptorSets[idx], 0, nullptr);
        VkBuffer vertexBuffers[] = {skyCubeBuffer->getRef()};
        VkDeviceSize offsets[] = {0};
        vkCmdBindVertexBuffers(graphicsBuffers[idx], 0, 1, vertexBuffers, offsets);
        vkCmdDraw(graphicsBuffers[idx], skyCubeBuffer->size, 1, 0, 0);
    }
    vkCmdEndRenderPass(graphicsBuffers[idx]);
}
void CommandCenter::recordUI(uint32_t idx, VkRenderPass uiPass, VkExtent2D extent, const std::function<void(VkCommandBuffer)>& call) {
    TracyVkZone(tracyCtx, graphicsBuffers[idx], "UI");

    VkRenderPassBeginInfo info = {};
    info.sType = VK_STRUCTURE_TYPE_RENDER_PASS_BEGIN_INFO;
    info.renderPass = uiPass;
    info.framebuffer = uiFramebuffers[idx];
    info.renderArea.extent = extent;
    info.clearValueCount = 1;
    VkClearValue clearValue{};
    clearValue.color = {0.f, 0.f, 0.f, 1.f};
    info.pClearValues = &clearValue;
    vkCmdBeginRenderPass(graphicsBuffers[idx], &info, VK_SUBPASS_CONTENTS_INLINE);

    call(graphicsBuffers[idx]);
    vkCmdEndRenderPass(graphicsBuffers[idx]);
}
void CommandCenter::submitGraphics(uint32_t idx, VkSemaphore waitSemaphore, VkSemaphore signalSemaphore, VkFence submittedFence) {
    assert(!freed);

    if (vkEndCommandBuffer(graphicsBuffers[idx]) != VK_SUCCESS) {
        FATAL("Failed to record graphics buffer!");
    }

    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;

    VkSemaphore waitSemaphores[] = {waitSemaphore};
    VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};
    submitInfo.waitSemaphoreCount = 1;
    submitInfo.pWaitSemaphores = waitSemaphores;
    submitInfo.pWaitDstStageMask = waitStages;

    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &graphicsBuffers[idx];

    VkSemaphore signalSemaphores[] = {signalSemaphore};
    submitInfo.signalSemaphoreCount = 1;
    submitInfo.pSignalSemaphores = signalSemaphores;

    vkResetFences(device, 1, &submittedFence);
    if (vkQueueSubmit(graphicsQueue, 1, &submitInfo, submittedFence) != VK_SUCCESS) {
        FATAL("Failed to submit draw command buffer!");
    }
}
